<commit_before>// Copyright 2019 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::fmt;
use std::io;
use std::mem;
use std::sync::{Arc, Mutex};
use std::thread;

use base::{error, warn, AsRawDescriptor, Error as SysError, Event, RawDescriptor};
use data_model::{DataInit, Le32};
use msg_socket::{MsgReceiver, MsgSender};
use resources::Alloc;
use vm_control::{FsMappingRequest, FsMappingRequestSocket, VmResponse};
use vm_memory::GuestMemory;

use crate::pci::{
    PciAddress, PciBarConfiguration, PciBarPrefetchable, PciBarRegionType, PciCapability,
};
use crate::virtio::{
    copy_config, DescriptorError, Interrupt, PciCapabilityType, Queue, VirtioDevice,
    VirtioPciShmCap, TYPE_FS,
};

mod multikey;
pub mod passthrough;
mod read_dir;
mod worker;

use fuse::Server;
use passthrough::PassthroughFs;
use worker::Worker;

// The fs device does not have a fixed number of queues.
const QUEUE_SIZE: u16 = 1024;

const FS_BAR_NUM: u8 = 4;
const FS_BAR_OFFSET: u64 = 0;
const FS_BAR_SIZE: u64 = 1 << 33;

/// Defined in kernel/include/uapi/linux/virtio_fs.h.
const VIRTIO_FS_SHMCAP_ID_CACHE: u8 = 0;

/// The maximum allowable length of the tag used to identify a specific virtio-fs device.
pub const FS_MAX_TAG_LEN: usize = 36;

/// kernel/include/uapi/linux/virtio_fs.h
#[repr(C, packed)]
#[derive(Clone, Copy)]
struct Config {
    /// Filesystem name (UTF-8, not NUL-terminated, padded with NULs)
    tag: [u8; FS_MAX_TAG_LEN],
    /// Number of request queues
    num_queues: Le32,
}

// Safe because all members are plain old data and any value is valid.
unsafe impl DataInit for Config {}

/// Errors that may occur during the creation or operation of an Fs device.
#[derive(Debug)]
pub enum Error {
    /// The tag for the Fs device was too long to fit in the config space.
    TagTooLong(usize),
    /// Failed to create the file system.
    CreateFs(io::Error),
    /// Creating WaitContext failed.
    CreateWaitContext(SysError),
    /// Error while polling for events.
    WaitError(SysError),
    /// Error while reading from the virtio queue's Event.
    ReadQueueEvent(SysError),
    /// A request is missing readable descriptors.
    NoReadableDescriptors,
    /// A request is missing writable descriptors.
    NoWritableDescriptors,
    /// Failed to signal the virio used queue.
    SignalUsedQueue(SysError),
    /// The `len` field of the header is too small.
    InvalidDescriptorChain(DescriptorError),
    /// Error happened in FUSE.
    FuseError(fuse::Error),
}

impl ::std::error::Error for Error {}

impl From<fuse::Error> for Error {
    fn from(err: fuse::Error) -> Error {
        Error::FuseError(err)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use Error::*;
        match self {
            TagTooLong(len) => write!(
                f,
                "Fs device tag is too long: len = {}, max = {}",
                len, FS_MAX_TAG_LEN
            ),
            CreateFs(err) => write!(f, "failed to create file system: {}", err),
            CreateWaitContext(err) => write!(f, "failed to create WaitContext: {}", err),
            WaitError(err) => write!(f, "failed to wait for events: {}", err),
            ReadQueueEvent(err) => write!(f, "failed to read from virtio queue Event: {}", err),
            NoReadableDescriptors => write!(f, "request does not have any readable descriptors"),
            NoWritableDescriptors => write!(f, "request does not have any writable descriptors"),
            SignalUsedQueue(err) => write!(f, "failed to signal used queue: {}", err),
            InvalidDescriptorChain(err) => write!(f, "DescriptorChain is invalid: {}", err),
            FuseError(err) => write!(f, "fuse error: {}", err),
        }
    }
}

pub type Result<T> = ::std::result::Result<T, Error>;

pub struct Fs {
    cfg: Config,
    fs: Option<PassthroughFs>,
    queue_sizes: Box<[u16]>,
    avail_features: u64,
    acked_features: u64,
    pci_bar: Option<Alloc>,
    socket: Option<FsMappingRequestSocket>,
    workers: Vec<(Event, thread::JoinHandle<Result<()>>)>,
}

impl Fs {
    pub fn new(
        base_features: u64,
        tag: &str,
        num_workers: usize,
        fs_cfg: passthrough::Config,
        socket: FsMappingRequestSocket,
    ) -> Result<Fs> {
        if tag.len() > FS_MAX_TAG_LEN {
            return Err(Error::TagTooLong(tag.len()));
        }

        let mut cfg_tag = [0u8; FS_MAX_TAG_LEN];
        cfg_tag[..tag.len()].copy_from_slice(tag.as_bytes());

        let cfg = Config {
            tag: cfg_tag,
            num_queues: Le32::from(num_workers as u32),
        };

        let fs = PassthroughFs::new(fs_cfg).map_err(Error::CreateFs)?;

        // There is always a high priority queue in addition to the request queues.
        let num_queues = num_workers + 1;

        Ok(Fs {
            cfg,
            fs: Some(fs),
            queue_sizes: vec![QUEUE_SIZE; num_queues].into_boxed_slice(),
            avail_features: base_features,
            acked_features: 0,
            pci_bar: None,
            socket: Some(socket),
            workers: Vec::with_capacity(num_workers + 1),
        })
    }

    fn stop_workers(&mut self) {
        for (kill_evt, handle) in mem::replace(&mut self.workers, Vec::new()) {
            if let Err(e) = kill_evt.write(1) {
                error!("failed to kill virtio-fs worker thread: {}", e);
                continue;
            }

            // Only wait on the child thread if we were able to send it a kill event.
            match handle.join() {
                Ok(r) => {
                    if let Err(e) = r {
                        error!("virtio-fs worker thread exited with error: {}", e)
                    }
                }
                Err(e) => error!("virtio-fs worker thread panicked: {:?}", e),
            }
        }
    }
}

impl VirtioDevice for Fs {
    fn keep_rds(&self) -> Vec<RawDescriptor> {
        let mut fds = self
            .fs
            .as_ref()
            .map(PassthroughFs::keep_rds)
            .unwrap_or_else(Vec::new);
        if let Some(rd) = self.socket.as_ref().map(|s| s.as_raw_descriptor()) {
            fds.push(rd);
        }

        fds
    }

    fn device_type(&self) -> u32 {
        TYPE_FS
    }

    fn queue_max_sizes(&self) -> &[u16] {
        &self.queue_sizes
    }

    fn features(&self) -> u64 {
        self.avail_features
    }

    fn ack_features(&mut self, mut v: u64) {
        // Check if the guest is ACK'ing a feature that we didn't claim to have.
        let unrequested_features = v & !self.avail_features;
        if unrequested_features != 0 {
            warn!("virtio_fs got unknown feature ack: {:x}", v);

            // Don't count these features as acked.
            v &= !unrequested_features;
        }
        self.acked_features |= v;
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        copy_config(data, 0, self.cfg.as_slice(), offset)
    }

    fn activate(
        &mut self,
        guest_mem: GuestMemory,
        interrupt: Interrupt,
        queues: Vec<Queue>,
        queue_evts: Vec<Event>,
    ) {
        if queues.len() != self.queue_sizes.len() || queue_evts.len() != self.queue_sizes.len() {
            return;
        }

        let fs = self.fs.take().expect("missing file system implementation");

        let server = Arc::new(Server::new(fs));
        let irq = Arc::new(interrupt);
        let socket = self.socket.take().expect("missing mapping socket");

        // Create the shared memory region now before we start processing requests.
        let request = FsMappingRequest::AllocateSharedMemoryRegion(
            self.pci_bar.as_ref().cloned().expect("No pci_bar"),
        );
        socket
            .send(&request)
            .expect("failed to send allocation message");
        let slot = match socket.recv() {
            Ok(VmResponse::RegisterMemory { pfn: _, slot }) => slot,
            Ok(VmResponse::Err(e)) => panic!("failed to allocate shared memory region: {}", e),
            r => panic!(
                "unexpected response to allocate shared memory region: {:?}",
                r
            ),
        };
        let socket = Arc::new(Mutex::new(socket));
        let mut watch_resample_event = true;
        for (idx, (queue, evt)) in queues.into_iter().zip(queue_evts.into_iter()).enumerate() {
            let (self_kill_evt, kill_evt) = match Event::new().and_then(|e| Ok((e.try_clone()?, e)))
            {
                Ok(v) => v,
                Err(e) => {
                    error!("fs: failed creating kill Event pair: {}", e);
                    self.stop_workers();
                    return;
                }
            };

            let mem = guest_mem.clone();
            let server = server.clone();
            let irq = irq.clone();
            let socket = Arc::clone(&socket);

            let worker_result = thread::Builder::new()
                .name(format!("virtio-fs worker {}", idx))
                .spawn(move || {
                    let mut worker = Worker::new(mem, queue, server, irq, socket, slot);
                    worker.run(evt, kill_evt, watch_resample_event)
                });

            if watch_resample_event {
                watch_resample_event = false;
            }

            match worker_result {
                Ok(worker) => self.workers.push((self_kill_evt, worker)),
                Err(e) => {
                    error!("fs: failed to spawn virtio_fs worker: {}", e);
                    self.stop_workers();
                    return;
                }
            }
        }
    }

    fn get_device_bars(&mut self, address: PciAddress) -> Vec<PciBarConfiguration> {
        self.pci_bar = Some(Alloc::PciBar {
            bus: address.bus,
            dev: address.dev,
            func: address.func,
            bar: FS_BAR_NUM,
        });

        vec![PciBarConfiguration::new(
            FS_BAR_NUM as usize,
            FS_BAR_SIZE,
            PciBarRegionType::Memory64BitRegion,
            PciBarPrefetchable::NotPrefetchable,
        )]
    }

    fn get_device_caps(&self) -> Vec<Box<dyn PciCapability>> {
        vec![Box::new(VirtioPciShmCap::new(
            PciCapabilityType::SharedMemoryConfig,
            FS_BAR_NUM,
            FS_BAR_OFFSET,
            FS_BAR_SIZE,
            VIRTIO_FS_SHMCAP_ID_CACHE,
        ))]
    }
}

impl Drop for Fs {
    fn drop(&mut self) {
        self.stop_workers()
    }
}
<commit_msg>devices: virtio: fs: Avoid allocating shared memory on ARM<commit_after>// Copyright 2019 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::fmt;
use std::io;
use std::mem;
use std::sync::{Arc, Mutex};
use std::thread;

use base::{error, warn, AsRawDescriptor, Error as SysError, Event, RawDescriptor};
use data_model::{DataInit, Le32};
use msg_socket::{MsgReceiver, MsgSender};
use resources::Alloc;
use vm_control::{FsMappingRequest, FsMappingRequestSocket, VmResponse};
use vm_memory::GuestMemory;

use crate::pci::{
    PciAddress, PciBarConfiguration, PciBarPrefetchable, PciBarRegionType, PciCapability,
};
use crate::virtio::{
    copy_config, DescriptorError, Interrupt, PciCapabilityType, Queue, VirtioDevice,
    VirtioPciShmCap, TYPE_FS,
};

mod multikey;
pub mod passthrough;
mod read_dir;
mod worker;

use fuse::Server;
use passthrough::PassthroughFs;
use worker::Worker;

// The fs device does not have a fixed number of queues.
const QUEUE_SIZE: u16 = 1024;

const FS_BAR_NUM: u8 = 4;
const FS_BAR_OFFSET: u64 = 0;
const FS_BAR_SIZE: u64 = 1 << 33;

/// Defined in kernel/include/uapi/linux/virtio_fs.h.
const VIRTIO_FS_SHMCAP_ID_CACHE: u8 = 0;

/// The maximum allowable length of the tag used to identify a specific virtio-fs device.
pub const FS_MAX_TAG_LEN: usize = 36;

/// kernel/include/uapi/linux/virtio_fs.h
#[repr(C, packed)]
#[derive(Clone, Copy)]
struct Config {
    /// Filesystem name (UTF-8, not NUL-terminated, padded with NULs)
    tag: [u8; FS_MAX_TAG_LEN],
    /// Number of request queues
    num_queues: Le32,
}

// Safe because all members are plain old data and any value is valid.
unsafe impl DataInit for Config {}

/// Errors that may occur during the creation or operation of an Fs device.
#[derive(Debug)]
pub enum Error {
    /// The tag for the Fs device was too long to fit in the config space.
    TagTooLong(usize),
    /// Failed to create the file system.
    CreateFs(io::Error),
    /// Creating WaitContext failed.
    CreateWaitContext(SysError),
    /// Error while polling for events.
    WaitError(SysError),
    /// Error while reading from the virtio queue's Event.
    ReadQueueEvent(SysError),
    /// A request is missing readable descriptors.
    NoReadableDescriptors,
    /// A request is missing writable descriptors.
    NoWritableDescriptors,
    /// Failed to signal the virio used queue.
    SignalUsedQueue(SysError),
    /// The `len` field of the header is too small.
    InvalidDescriptorChain(DescriptorError),
    /// Error happened in FUSE.
    FuseError(fuse::Error),
}

impl ::std::error::Error for Error {}

impl From<fuse::Error> for Error {
    fn from(err: fuse::Error) -> Error {
        Error::FuseError(err)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use Error::*;
        match self {
            TagTooLong(len) => write!(
                f,
                "Fs device tag is too long: len = {}, max = {}",
                len, FS_MAX_TAG_LEN
            ),
            CreateFs(err) => write!(f, "failed to create file system: {}", err),
            CreateWaitContext(err) => write!(f, "failed to create WaitContext: {}", err),
            WaitError(err) => write!(f, "failed to wait for events: {}", err),
            ReadQueueEvent(err) => write!(f, "failed to read from virtio queue Event: {}", err),
            NoReadableDescriptors => write!(f, "request does not have any readable descriptors"),
            NoWritableDescriptors => write!(f, "request does not have any writable descriptors"),
            SignalUsedQueue(err) => write!(f, "failed to signal used queue: {}", err),
            InvalidDescriptorChain(err) => write!(f, "DescriptorChain is invalid: {}", err),
            FuseError(err) => write!(f, "fuse error: {}", err),
        }
    }
}

pub type Result<T> = ::std::result::Result<T, Error>;

pub struct Fs {
    cfg: Config,
    fs: Option<PassthroughFs>,
    queue_sizes: Box<[u16]>,
    avail_features: u64,
    acked_features: u64,
    pci_bar: Option<Alloc>,
    socket: Option<FsMappingRequestSocket>,
    workers: Vec<(Event, thread::JoinHandle<Result<()>>)>,
}

impl Fs {
    pub fn new(
        base_features: u64,
        tag: &str,
        num_workers: usize,
        fs_cfg: passthrough::Config,
        socket: FsMappingRequestSocket,
    ) -> Result<Fs> {
        if tag.len() > FS_MAX_TAG_LEN {
            return Err(Error::TagTooLong(tag.len()));
        }

        let mut cfg_tag = [0u8; FS_MAX_TAG_LEN];
        cfg_tag[..tag.len()].copy_from_slice(tag.as_bytes());

        let cfg = Config {
            tag: cfg_tag,
            num_queues: Le32::from(num_workers as u32),
        };

        let fs = PassthroughFs::new(fs_cfg).map_err(Error::CreateFs)?;

        // There is always a high priority queue in addition to the request queues.
        let num_queues = num_workers + 1;

        Ok(Fs {
            cfg,
            fs: Some(fs),
            queue_sizes: vec![QUEUE_SIZE; num_queues].into_boxed_slice(),
            avail_features: base_features,
            acked_features: 0,
            pci_bar: None,
            socket: Some(socket),
            workers: Vec::with_capacity(num_workers + 1),
        })
    }

    fn stop_workers(&mut self) {
        for (kill_evt, handle) in mem::replace(&mut self.workers, Vec::new()) {
            if let Err(e) = kill_evt.write(1) {
                error!("failed to kill virtio-fs worker thread: {}", e);
                continue;
            }

            // Only wait on the child thread if we were able to send it a kill event.
            match handle.join() {
                Ok(r) => {
                    if let Err(e) = r {
                        error!("virtio-fs worker thread exited with error: {}", e)
                    }
                }
                Err(e) => error!("virtio-fs worker thread panicked: {:?}", e),
            }
        }
    }
}

impl VirtioDevice for Fs {
    fn keep_rds(&self) -> Vec<RawDescriptor> {
        let mut fds = self
            .fs
            .as_ref()
            .map(PassthroughFs::keep_rds)
            .unwrap_or_else(Vec::new);
        if let Some(rd) = self.socket.as_ref().map(|s| s.as_raw_descriptor()) {
            fds.push(rd);
        }

        fds
    }

    fn device_type(&self) -> u32 {
        TYPE_FS
    }

    fn queue_max_sizes(&self) -> &[u16] {
        &self.queue_sizes
    }

    fn features(&self) -> u64 {
        self.avail_features
    }

    fn ack_features(&mut self, mut v: u64) {
        // Check if the guest is ACK'ing a feature that we didn't claim to have.
        let unrequested_features = v & !self.avail_features;
        if unrequested_features != 0 {
            warn!("virtio_fs got unknown feature ack: {:x}", v);

            // Don't count these features as acked.
            v &= !unrequested_features;
        }
        self.acked_features |= v;
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        copy_config(data, 0, self.cfg.as_slice(), offset)
    }

    fn activate(
        &mut self,
        guest_mem: GuestMemory,
        interrupt: Interrupt,
        queues: Vec<Queue>,
        queue_evts: Vec<Event>,
    ) {
        if queues.len() != self.queue_sizes.len() || queue_evts.len() != self.queue_sizes.len() {
            return;
        }

        let fs = self.fs.take().expect("missing file system implementation");

        let server = Arc::new(Server::new(fs));
        let irq = Arc::new(interrupt);
        let socket = self.socket.take().expect("missing mapping socket");
        let mut slot = 0;

        // Set up shared memory for DAX.
        // TODO(b/176129399): Remove cfg! once DAX is supported on ARM.
        if cfg!(any(target_arch = "x86", target_arch = "x86_64")) {
            // Create the shared memory region now before we start processing requests.
            let request = FsMappingRequest::AllocateSharedMemoryRegion(
                self.pci_bar.as_ref().cloned().expect("No pci_bar"),
            );
            socket
                .send(&request)
                .expect("failed to send allocation message");
            slot = match socket.recv() {
                Ok(VmResponse::RegisterMemory { pfn: _, slot }) => slot,
                Ok(VmResponse::Err(e)) => panic!("failed to allocate shared memory region: {}", e),
                r => panic!(
                    "unexpected response to allocate shared memory region: {:?}",
                    r
                ),
            };
        }

        let socket = Arc::new(Mutex::new(socket));
        let mut watch_resample_event = true;
        for (idx, (queue, evt)) in queues.into_iter().zip(queue_evts.into_iter()).enumerate() {
            let (self_kill_evt, kill_evt) = match Event::new().and_then(|e| Ok((e.try_clone()?, e)))
            {
                Ok(v) => v,
                Err(e) => {
                    error!("fs: failed creating kill Event pair: {}", e);
                    self.stop_workers();
                    return;
                }
            };

            let mem = guest_mem.clone();
            let server = server.clone();
            let irq = irq.clone();
            let socket = Arc::clone(&socket);

            let worker_result = thread::Builder::new()
                .name(format!("virtio-fs worker {}", idx))
                .spawn(move || {
                    let mut worker = Worker::new(mem, queue, server, irq, socket, slot);
                    worker.run(evt, kill_evt, watch_resample_event)
                });

            if watch_resample_event {
                watch_resample_event = false;
            }

            match worker_result {
                Ok(worker) => self.workers.push((self_kill_evt, worker)),
                Err(e) => {
                    error!("fs: failed to spawn virtio_fs worker: {}", e);
                    self.stop_workers();
                    return;
                }
            }
        }
    }

    fn get_device_bars(&mut self, address: PciAddress) -> Vec<PciBarConfiguration> {
        self.pci_bar = Some(Alloc::PciBar {
            bus: address.bus,
            dev: address.dev,
            func: address.func,
            bar: FS_BAR_NUM,
        });

        vec![PciBarConfiguration::new(
            FS_BAR_NUM as usize,
            FS_BAR_SIZE,
            PciBarRegionType::Memory64BitRegion,
            PciBarPrefetchable::NotPrefetchable,
        )]
    }

    fn get_device_caps(&self) -> Vec<Box<dyn PciCapability>> {
        vec![Box::new(VirtioPciShmCap::new(
            PciCapabilityType::SharedMemoryConfig,
            FS_BAR_NUM,
            FS_BAR_OFFSET,
            FS_BAR_SIZE,
            VIRTIO_FS_SHMCAP_ID_CACHE,
        ))]
    }
}

impl Drop for Fs {
    fn drop(&mut self) {
        self.stop_workers()
    }
}
<|endoftext|>


<commit_before>extern crate byteorder;

use self::byteorder::{ReadBytesExt, WriteBytesExt, BigEndian};
use std::error::Error;
use std::io::Write;

use backend::Debug;
use pg::{Pg, PgTypeMetadata};
use query_source::Queryable;
use types::{HasSqlType, FromSql, ToSql, Array, IsNull, NotNull, FromSqlRow};

impl<T> HasSqlType<Array<T>> for Pg where
    Pg: HasSqlType<T>,
{
    fn metadata() -> PgTypeMetadata {
        PgTypeMetadata {
            oid: <Pg as HasSqlType<T>>::metadata().array_oid,
            array_oid: 0,
        }
    }
}

impl<T> HasSqlType<Array<T>> for Debug where
    Debug: HasSqlType<T>,
{
    fn metadata() {}
}

impl_query_id!(Array<T>);

impl<T> NotNull for Array<T> {
}

impl<T, ST> FromSql<Array<ST>, Pg> for Vec<T> where
    T: FromSql<ST, Pg>,
    Pg: HasSqlType<ST>,
{
    fn from_sql(bytes: Option<&[u8]>) -> Result<Self, Box<Error+Send+Sync>> {
        let mut bytes = not_none!(bytes);
        let num_dimensions = try!(bytes.read_i32::<BigEndian>());
        let has_null = try!(bytes.read_i32::<BigEndian>()) != 0;
        let _oid = try!(bytes.read_i32::<BigEndian>());

        if num_dimensions == 0 {
            return Ok(Vec::new())
        }

        let num_elements = try!(bytes.read_i32::<BigEndian>());
        let lower_bound = try!(bytes.read_i32::<BigEndian>());

        assert!(num_dimensions == 1, "multi-dimensional arrays are not supported");
        assert!(lower_bound == 1, "lower bound must be 1");

        (0..num_elements).map(|_| {
            let elem_size = try!(bytes.read_i32::<BigEndian>());
            if has_null && elem_size == -1 {
                T::from_sql(None)
            } else {
                let (elem_bytes, new_bytes) = bytes.split_at(elem_size as usize);
                bytes = new_bytes;
                T::from_sql(Some(&elem_bytes))
            }
        }).collect()
    }
}

impl<T, ST> FromSqlRow<Array<ST>, Pg> for Vec<T> where
    Pg: HasSqlType<ST>,
    Vec<T>: FromSql<Array<ST>, Pg>,
{
    fn build_from_row<R: ::row::Row<Pg>>(row: &mut R) -> Result<Self, Box<Error+Send+Sync>> {
        FromSql::<Array<ST>, Pg>::from_sql(row.take())
    }
}

impl<T, ST> Queryable<Array<ST>, Pg> for Vec<T> where
    T: FromSql<ST, Pg> + Queryable<ST, Pg>,
    Pg: HasSqlType<ST>,
{
    type Row = Self;
    fn build(row: Self) -> Self {
        row
    }
}

use expression::AsExpression;
use expression::bound::Bound;

impl<'a, ST, T> AsExpression<Array<ST>> for &'a [T] where
    Pg: HasSqlType<ST>,
{
    type Expression = Bound<Array<ST>, Self>;

    fn as_expression(self) -> Self::Expression {
        Bound::new(self)
    }
}

impl<ST, T> AsExpression<Array<ST>> for Vec<T> where
    Pg: HasSqlType<ST>,
{
    type Expression = Bound<Array<ST>, Self>;

    fn as_expression(self) -> Self::Expression {
        Bound::new(self)
    }
}

impl<'a, ST, T> AsExpression<Array<ST>> for &'a Vec<T> where
    Pg: HasSqlType<ST>,
{
    type Expression = Bound<Array<ST>, Self>;

    fn as_expression(self) -> Self::Expression {
        Bound::new(self)
    }
}

impl<'a, ST, T> ToSql<Array<ST>, Pg> for &'a [T] where
    Pg: HasSqlType<ST>,
    T: ToSql<ST, Pg>,
{
    fn to_sql<W: Write>(&self, out: &mut W) -> Result<IsNull, Box<Error+Send+Sync>> {
        let num_dimensions = 1;
        try!(out.write_i32::<BigEndian>(num_dimensions));
        let flags = 0;
        try!(out.write_i32::<BigEndian>(flags));
        try!(out.write_u32::<BigEndian>(Pg::metadata().oid));
        try!(out.write_i32::<BigEndian>(self.len() as i32));
        let lower_bound = 1;
        try!(out.write_i32::<BigEndian>(lower_bound));

        let mut buffer = Vec::new();
        for elem in self.iter() {
            let is_null = try!(elem.to_sql(&mut buffer));
            assert!(is_null == IsNull::No, "Arrays containing null are not supported");
            try!(out.write_i32::<BigEndian>(buffer.len() as i32));
            try!(out.write_all(&buffer));
            buffer.clear();
        }

        Ok(IsNull::No)
    }
}

impl<ST, T> ToSql<Array<ST>, Pg> for Vec<T> where
    Pg: HasSqlType<ST>,
    for<'a> &'a [T]: ToSql<Array<ST>, Pg>,
{
    fn to_sql<W: Write>(&self, out: &mut W) -> Result<IsNull, Box<Error+Send+Sync>> {
        (&self as &[T]).to_sql(out)
    }
}
<commit_msg>Implement `AsExpression<Nullable<Array<T>>` for various types<commit_after>extern crate byteorder;

use self::byteorder::{ReadBytesExt, WriteBytesExt, BigEndian};
use std::error::Error;
use std::io::Write;

use backend::Debug;
use pg::{Pg, PgTypeMetadata};
use query_source::Queryable;
use types::*;

impl<T> HasSqlType<Array<T>> for Pg where
    Pg: HasSqlType<T>,
{
    fn metadata() -> PgTypeMetadata {
        PgTypeMetadata {
            oid: <Pg as HasSqlType<T>>::metadata().array_oid,
            array_oid: 0,
        }
    }
}

impl<T> HasSqlType<Array<T>> for Debug where
    Debug: HasSqlType<T>,
{
    fn metadata() {}
}

impl_query_id!(Array<T>);

impl<T> NotNull for Array<T> {
}

impl<T, ST> FromSql<Array<ST>, Pg> for Vec<T> where
    T: FromSql<ST, Pg>,
    Pg: HasSqlType<ST>,
{
    fn from_sql(bytes: Option<&[u8]>) -> Result<Self, Box<Error+Send+Sync>> {
        let mut bytes = not_none!(bytes);
        let num_dimensions = try!(bytes.read_i32::<BigEndian>());
        let has_null = try!(bytes.read_i32::<BigEndian>()) != 0;
        let _oid = try!(bytes.read_i32::<BigEndian>());

        if num_dimensions == 0 {
            return Ok(Vec::new())
        }

        let num_elements = try!(bytes.read_i32::<BigEndian>());
        let lower_bound = try!(bytes.read_i32::<BigEndian>());

        assert!(num_dimensions == 1, "multi-dimensional arrays are not supported");
        assert!(lower_bound == 1, "lower bound must be 1");

        (0..num_elements).map(|_| {
            let elem_size = try!(bytes.read_i32::<BigEndian>());
            if has_null && elem_size == -1 {
                T::from_sql(None)
            } else {
                let (elem_bytes, new_bytes) = bytes.split_at(elem_size as usize);
                bytes = new_bytes;
                T::from_sql(Some(&elem_bytes))
            }
        }).collect()
    }
}

impl<T, ST> FromSqlRow<Array<ST>, Pg> for Vec<T> where
    Pg: HasSqlType<ST>,
    Vec<T>: FromSql<Array<ST>, Pg>,
{
    fn build_from_row<R: ::row::Row<Pg>>(row: &mut R) -> Result<Self, Box<Error+Send+Sync>> {
        FromSql::<Array<ST>, Pg>::from_sql(row.take())
    }
}

impl<T, ST> Queryable<Array<ST>, Pg> for Vec<T> where
    T: FromSql<ST, Pg> + Queryable<ST, Pg>,
    Pg: HasSqlType<ST>,
{
    type Row = Self;
    fn build(row: Self) -> Self {
        row
    }
}

use expression::AsExpression;
use expression::bound::Bound;

macro_rules! array_as_expression {
    ($ty:ty, $sql_type:ty) => {
        impl<'a, ST, T> AsExpression<$sql_type> for $ty where
            Pg: HasSqlType<ST>,
        {
            type Expression = Bound<$sql_type, Self>;

            fn as_expression(self) -> Self::Expression {
                Bound::new(self)
            }
        }
    }
}

array_as_expression!(&'a [T], Array<ST>);
array_as_expression!(&'a [T], Nullable<Array<ST>>);
array_as_expression!(Vec<T>, Array<ST>);
array_as_expression!(Vec<T>, Nullable<Array<ST>>);
array_as_expression!(&'a Vec<T>, Array<ST>);
array_as_expression!(&'a Vec<T>, Nullable<Array<ST>>);

impl<'a, ST, T> ToSql<Array<ST>, Pg> for &'a [T] where
    Pg: HasSqlType<ST>,
    T: ToSql<ST, Pg>,
{
    fn to_sql<W: Write>(&self, out: &mut W) -> Result<IsNull, Box<Error+Send+Sync>> {
        let num_dimensions = 1;
        try!(out.write_i32::<BigEndian>(num_dimensions));
        let flags = 0;
        try!(out.write_i32::<BigEndian>(flags));
        try!(out.write_u32::<BigEndian>(Pg::metadata().oid));
        try!(out.write_i32::<BigEndian>(self.len() as i32));
        let lower_bound = 1;
        try!(out.write_i32::<BigEndian>(lower_bound));

        let mut buffer = Vec::new();
        for elem in self.iter() {
            let is_null = try!(elem.to_sql(&mut buffer));
            assert!(is_null == IsNull::No, "Arrays containing null are not supported");
            try!(out.write_i32::<BigEndian>(buffer.len() as i32));
            try!(out.write_all(&buffer));
            buffer.clear();
        }

        Ok(IsNull::No)
    }
}

impl<ST, T> ToSql<Array<ST>, Pg> for Vec<T> where
    Pg: HasSqlType<ST>,
    for<'a> &'a [T]: ToSql<Array<ST>, Pg>,
{
    fn to_sql<W: Write>(&self, out: &mut W) -> Result<IsNull, Box<Error+Send+Sync>> {
        (&self as &[T]).to_sql(out)
    }
}
<|endoftext|>


<commit_before>#![no_std]
#![allow(ctypes)]

#![feature(lang_items)]
#[lang="sized"]
trait Sized {}

#[lang="copy"]
trait Copy {}

impl Copy for Color {}

enum Color {
    Black      = 0,
    Blue       = 1,
    Green      = 2,
    Cyan       = 3,
    Red        = 4,
    Pink       = 5,
    Brown      = 6,
    LightGray  = 7,
    DarkGray   = 8,
    LightBlue  = 9,
    LightGreen = 10,
    LightCyan  = 11,
    LightRed   = 12,
    LightPink  = 13,
    Yellow     = 14,
    White      = 15,
}

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

struct IntRange {
    cur: int,
    max: int
}

impl IntRange {
    fn next(&mut self) -> Option<int> {
        if self.cur < self.max {
            self.cur += 1;
            Option::Some(self.cur - 1)
        } else {
            Option::None
        }
    }
}

fn range(lo: int, hi: int) -> IntRange {
    IntRange { cur: lo, max: hi }
}

fn clear_screen(background: Color) {
    let mut r = range(0, 80 * 25);
    loop {
        match r.next() {
            Option::Some(x) => {
                unsafe {
                    *((0xb8000 + x * 2) as *mut u16) = (background as u16) << 12;
                }
            },
            Option::None =>{break}
        }
    }
}

#[no_mangle]
#[no_split_stack]
pub fn main() {
    clear_screen(Color::LightRed);
}
<commit_msg>No more warnings<commit_after>#![no_std]
#![allow(improper_ctypes)]

#![feature(lang_items)]
#[lang="sized"]
trait Sized {}

#[lang="copy"]
trait Copy {}

impl Copy for Color {}

#[allow(dead_code)]
enum Color {
    Black      = 0,
    Blue       = 1,
    Green      = 2,
    Cyan       = 3,
    Red        = 4,
    Pink       = 5,
    Brown      = 6,
    LightGray  = 7,
    DarkGray   = 8,
    LightBlue  = 9,
    LightGreen = 10,
    LightCyan  = 11,
    LightRed   = 12,
    LightPink  = 13,
    Yellow     = 14,
    White      = 15,
}

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

struct IntRange {
    cur: isize,
    max: isize
}

impl IntRange {
    fn next(&mut self) -> Option<isize> {
        if self.cur < self.max {
            self.cur += 1;
            Option::Some(self.cur - 1)
        } else {
            Option::None
        }
    }
}

fn range(lo: isize, hi: isize) -> IntRange {
    IntRange { cur: lo, max: hi }
}

fn clear_screen(background: Color) {
    let mut r = range(0, 80 * 25);
    loop {
        match r.next() {
            Option::Some(x) => {
                unsafe {
                    *((0xb8000 + x * 2) as *mut u16) = (background as u16) << 12;
                }
            },
            Option::None =>{break}
        }
    }
}

#[no_mangle]
#[no_stack_check]
pub fn main() {
    clear_screen(Color::LightRed);
}
<|endoftext|>


<commit_before>// library
extern crate my_extern_lib;

// modules
mod simple_functions {
    pub fn print_hello_world() {
        println!("Hello, World! (from the simple_functions module)");
     }

    // nested modules
    pub mod simple_math {
        pub fn multiple_two_integers(a: i32, b: i32) -> i32 {
            a * b
        }
    }
}

use simple_functions::simple_math::multiple_two_integers;

fn print_some_msg() {
    println!("Modules: printing from the outer function");
}

mod sample_module {

    pub fn call_sample_module() {
        use second_module::print_some_msg as print_second;
        use self::second_module::print_some_msg as print_second_sample;
        use super::print_some_msg as print_outside;
        print_outside();
        print_second();
        print_some_msg();
        print_second_sample();
    }

    mod second_module {

        pub fn print_some_msg() {
            println!("Modules: printing from the second module inside the sample module");
        }

    }

    fn print_some_msg() {
        println!("Modules: printing from the inside of the sample module");
    }

}

mod second_module {
    pub fn print_some_msg() {
        println!("Modules: printing from the second module");
    }
}

mod modules;

fn main() {

    // formatted output
    println!("Hello, Rust!");
    println!("Integer number = {}", 23);
    println!("Float number = {}", 1.012);
    println!("Repeat printing: {0}, {1}, {1}, {3}, {2}", 2, 1, 44, 9);
    println!("Named arguments: {name}, {age}, {occupation}", age="24", name="Ann", occupation="Designer");
    println!("{:08b} XOR {:08b} = {:08b}", 3i8, 32i8, 3i8 ^ 34i8);
    println!("Floating point number {:?}", 2.910210);
    println!("Hex: {:08X}, {:08x}", 131342112, 131342112);
    // TODO find out why this doesn't work:
    //println!("{0:08X}, {0:08x}", 7661);

    // variables
    let integer = 3i32;
    let floating_point = 2.1137;

    let mut u_integer = 3u32;
    //integer += 1; error: an immutable variable
    u_integer += 1;

    println!("{}, {}, {}", integer, floating_point, u_integer);
    // TODO find out why snake case instead of camel case

    // shadowing
    let first_var = 4;
    {
        let first_var = 51;
        println!("first_var = {}", first_var);
    }
    println!("first_var = {}", first_var);

    // unitialized variables
    let uninit_var;
    //println!("uninit_var = {}", uninit_var); error: use of an init variable
    uninit_var = 231.2;
    println!("already initialized uninit_var = {}", uninit_var);

    // type casting
    let a_integer = 0x40_u8;
    let a_char = a_integer as char;
    println!("char cast from integer {}: {}", a_integer, a_char);

    let a_float = 8.45511;
    println!("integer cast from float {}: {}", a_float, a_float as i32);

    println!("Binary represntation of {}: {:08b}", -1_i8, -1_i8);
    println!("Binary represntation of {}: {:08b}", -2_i8, -2_i8);

    // aliases
    type FloatingPoint32 = f32;
    let a_float_2 : FloatingPoint32 = 3 as f32;
    println!("Type alias: {}, size of var {}", a_float_2, std::mem::size_of_val(&a_float_2));

    // expressions
    let expr = {
        let x = 4.09;
        let y = 1.28;
        println!("x = {}, y = {}", x, y);
        x / y
    };
    println!("expr = {}", expr);

    // if-statement (every if-statement is an expression too, which is really a cool thing)
    let speed = 10;
    let mass = 890;
    let consumption =
        if speed < 0 {
            -1.0
        }
        else if speed == 0 {
            0.8
        }
        else if speed > 0 && speed < 10 {
            18.52
        }
        else if speed >= 10 && speed < 30 {
            let x = mass / speed;
            (x as f32) * 0.05
        }
        else {
            11.71
        };
        println!("Speed = {}, consumption = {}", speed, consumption);

        // loop
        let mut count = 0u32;
        loop {
            println!("loop iteration #{}", count);
            count += 1;
            if count > 9 {
            break;
            }
        }

        count = 0;
        while count < 10 {
            println!("while loop iteration #{}", count);
            count += 1;
        }

        for n in 0u32..10 {
            println!("for loop iteration #{}", n);
        }

        // matching
        let age = 19;
        match age {
            1  => println!("Seems like 1"),
            5  => println!("It's 5"),
            18 => println!("Exactly 18!"),
            _  => println!("I couldn't guess :("),
        }

        match age {
            a if a > 1 && a <= 9  => println!("The age is in range [1,9]"),
            a if a > 9 && a <= 19 => println!("The age is in range [10, 19]"),
            _                     => println!("Nothing, I don't know what range the age is in :("),
        }

        // the previous match is equivalent to the following:
        match age {
            a @ 1...9   => println!("The age is {} and in range [1,9]", a),
            a @ 10...19 => println!("The age is {} and in range [10, 19]", a),
            a           => println!("Nothing, I don't know what range the age {} is in :(", a),
        }

        // if let
        let a_some = Some(0.0008);
        if let Some(i) = a_some {
            println!("a_some is not None, {}", i);
            println!("result {}",  9.0 / i);
        }
        else {
            println!("None");
        }

        // while let
        let mut a_optional = Some(10);

        while let Some(i) = a_optional {
            if i == 20 {
                println!("while let: i equals 20, leave");
                a_optional = None;
            }
            else {
                a_optional = Some(i + 1);
            }
        }

        // functions, finally!
        let first_integer: i32 = -20;
        let second_integer: i32 = 8;
        let result_of_addition = add_two_integers(first_integer, second_integer);
        println!("{} + {} = {}", first_integer, second_integer, result_of_addition);

        simple_functions::print_hello_world();
        let mul_result = simple_functions::simple_math::multiple_two_integers(3, 15);
        println!("Modules: multiplication result = {}", mul_result);

        // this doesn't work.. why?
        println!("Modules: multiplication result 2 = {}", multiple_two_integers(-1, 8));

        // "super" and "self"
        sample_module::call_sample_module();

        // file hierarchy
        {
            let a : i32 = 4;
            let b : i32 = 15;
            println!("modules::math {} + {} = {}", a, b, modules::math::add_two_int(a, b));
        }

        // call external library
        my_extern_lib::run(5);

        // tuples
        let pair = ("key", 4);
        println!("tuple value: ({}, {})", pair.0, pair.1);
        println!("tuple value: {:?}", pair);

        // structs
        let bmw: Car = Car {color: "Blue".to_string(), weight: 940};
        println!("BMW: {} - {}", bmw.color, bmw.weight);

        // tuple structs
        let color_1: RgbColor = RgbColor(123, 33, 81);
        println!("RGB color: {}, {}, {}", color_1.0, color_1.1, color_1.2);

        // unit structs
        let _nil = Nil;

        // a struct with a constructor and getter
        let rabbit = my_extern_lib::Animal::new(2, 4, "The King");
        println!("Rabbit: {}, {}, {}", rabbit.age, rabbit.weight, rabbit.get_secret_info());

        // generics
        {
            struct ThreeValues<T, U, V> {
                first: T,
                second: U,
                third: V,
            }

            let three_values = ThreeValues { first: 4, second: 2.01, third: "hello" };
            println!("Generic data structure: {}, {}, {}", three_values.first, three_values.second, three_values.third);

            // generic function
            fn swap<U, V>(first_value: U, second_value: V) -> (V, U) {
                (second_value, first_value)
            }

            let random_tuple = swap(4, 2);
            println!("Use unused variables: {} {}", random_tuple.0, random_tuple.1);

        }

        // box (heap allocation)
        {

            let toyota: Car = Car {color: "Green".to_string(), weight: 1040,};
            let porche: Box<Car> = Box::new(Car {color: "Black".to_string(), weight: 876, });

            println!("Toyota's size on stack is {}", std::mem::size_of_val(&toyota));
            println!("Boxed porche's size on stack is {}", std::mem::size_of_val(&porche));

            let unboxed_porche: Car = *porche;
            println!("Unboxed porche's size on stack is {}", std::mem::size_of_val(&unboxed_porche));

            // ownership
            let bmw : Box<Car> = Box::new(Car {color: "Yellow".to_string(), weight: 1620});
            println!("Boxed BMW at {:p}: {}/{}", &bmw, bmw.color, bmw.weight);
            let moved_bmw: Box<Car> = bmw;
            println!("Moved BMW at {:p}: {}/{}", &moved_bmw, moved_bmw.color, moved_bmw.weight);
            // can't do the following, the ownership has been moved
            // println!("Boxed BMW at {:p}: {}/{}", &bmw, bmw.color, bmw.weight);

            // mutability
            let immutable_car = Box::new(Car {color: "Pink".to_string(), weight: 510});
            // can't do this
            // immutable_car.weight = 403;

            // but we can change mutability
            let mut mutable_car = immutable_car;
            mutable_car.weight = 401;

        }

        // ref
        {

            let mut bmw: Car = Car {color: "Green".to_string(), weight: 870};
            println!("Ref: bmw {},{}", bmw.color, bmw.weight);

            {
                let Car {color: _, weight: ref weight_reference} = bmw;

                println!("Ref: {}", *weight_reference);
            }

            {
                let Car {color: _, weight: ref mut weight_mut_ref} = bmw;

                *weight_mut_ref = 200;
            }

            println!("Ref: bmw {},{}", bmw.color, bmw.weight);

        }

        // try some enums
        {

            enum TwoValues {
                First(i8),
                Second(f32),
            }

            fn test_enum(some_number: i32) -> TwoValues {

                if some_number > 0 {
                    TwoValues::First(120)
                }
                else {
                    TwoValues::Second(3.09)
                }

            }

            match test_enum(-2) {
                TwoValues::First(value) => println!("TwoValues enum: {}", value),
                TwoValues::Second(value) => println!("TwoValues enum: {}", value),
            }

            let two_values_instance = TwoValues::First(9);
            println!("TwoValues enum: size of a variant {}", std::mem::size_of_val(&two_values_instance));

        }

        // closures
        {
            let say = |phrase: String| {println!("{}", phrase)};
            say("Hello, closure!".to_string()); // this is probably worse than just println!

            fn mutate_integer<F: Fn(i32) -> i32>(value: i32, mutator: F) -> i32 {
                mutator(value)
            }

            let modulo_three = |v: i32| {v % 3};

            let value: i32 = 23;
            println!("Closures: input value = {}, result = {}", value, mutate_integer(value, modulo_three));

            let add_five = |v: u32| {v + 5};
            let mul_ten = |v: u32| {v * 10};

            fn mutate_integer_2<F, G>(value: u32, mutator_1: F, mutator_2: G) -> u32
                where F: Fn(u32) -> u32, G: Fn(u32) -> u32 {

                mutator_2(mutator_1(value))

            }

            let value_2: u32 = 7194682;
            println!("Closures #2: input value = {}, result = {}", value_2, mutate_integer_2(value_2, add_five, mul_ten));
        }

        // threads
        {
            use std::thread;
            use std::old_io::timer;
            use std::time::Duration;

            // without a guard
            thread::spawn(|| {
                println!("Very first thread EEEEEEYYYYYEEEECCCCAAAAAAATTTCCCHHHEEEERRR!!!");
            });
            timer::sleep(Duration::milliseconds(50)); // we need it because main terminates and kills all the threads

            // with a guard
            let thread_guard = thread::scoped(|| {
                println!("Thread with a guard here!!");
            });

            // another example
            for i in range(0, 16) {

                thread::spawn(move || {
                    println!("Thread #{} started", i);
                    timer::sleep(Duration::milliseconds(500 % (i + 1)));
                    println!("Thread #{} finished", i);
                });

            }
            timer::sleep(Duration::milliseconds(2500));
        }
}

// a C-like struct
struct Car {
    color: String,
    weight: u32,
}

// a tuple struct s
struct RgbColor(i8, i8, i8);

// a unit struct
struct Nil;

fn add_two_integers(a: i32, b: i32) -> i32 {
    a + b
}
<commit_msg>Progress: channels<commit_after>// library
extern crate my_extern_lib;

// modules
mod simple_functions {
    pub fn print_hello_world() {
        println!("Hello, World! (from the simple_functions module)");
     }

    // nested modules
    pub mod simple_math {
        pub fn multiple_two_integers(a: i32, b: i32) -> i32 {
            a * b
        }
    }
}

use simple_functions::simple_math::multiple_two_integers;

fn print_some_msg() {
    println!("Modules: printing from the outer function");
}

mod sample_module {

    pub fn call_sample_module() {
        use second_module::print_some_msg as print_second;
        use self::second_module::print_some_msg as print_second_sample;
        use super::print_some_msg as print_outside;
        print_outside();
        print_second();
        print_some_msg();
        print_second_sample();
    }

    mod second_module {

        pub fn print_some_msg() {
            println!("Modules: printing from the second module inside the sample module");
        }

    }

    fn print_some_msg() {
        println!("Modules: printing from the inside of the sample module");
    }

}

mod second_module {
    pub fn print_some_msg() {
        println!("Modules: printing from the second module");
    }
}

mod modules;

fn main() {

    // formatted output
    println!("Hello, Rust!");
    println!("Integer number = {}", 23);
    println!("Float number = {}", 1.012);
    println!("Repeat printing: {0}, {1}, {1}, {3}, {2}", 2, 1, 44, 9);
    println!("Named arguments: {name}, {age}, {occupation}", age="24", name="Ann", occupation="Designer");
    println!("{:08b} XOR {:08b} = {:08b}", 3i8, 32i8, 3i8 ^ 34i8);
    println!("Floating point number {:?}", 2.910210);
    println!("Hex: {:08X}, {:08x}", 131342112, 131342112);
    // TODO find out why this doesn't work:
    //println!("{0:08X}, {0:08x}", 7661);

    // variables
    let integer = 3i32;
    let floating_point = 2.1137;

    let mut u_integer = 3u32;
    //integer += 1; error: an immutable variable
    u_integer += 1;

    println!("{}, {}, {}", integer, floating_point, u_integer);
    // TODO find out why snake case instead of camel case

    // shadowing
    let first_var = 4;
    {
        let first_var = 51;
        println!("first_var = {}", first_var);
    }
    println!("first_var = {}", first_var);

    // unitialized variables
    let uninit_var;
    //println!("uninit_var = {}", uninit_var); error: use of an init variable
    uninit_var = 231.2;
    println!("already initialized uninit_var = {}", uninit_var);

    // type casting
    let a_integer = 0x40_u8;
    let a_char = a_integer as char;
    println!("char cast from integer {}: {}", a_integer, a_char);

    let a_float = 8.45511;
    println!("integer cast from float {}: {}", a_float, a_float as i32);

    println!("Binary represntation of {}: {:08b}", -1_i8, -1_i8);
    println!("Binary represntation of {}: {:08b}", -2_i8, -2_i8);

    // aliases
    type FloatingPoint32 = f32;
    let a_float_2 : FloatingPoint32 = 3 as f32;
    println!("Type alias: {}, size of var {}", a_float_2, std::mem::size_of_val(&a_float_2));

    // expressions
    let expr = {
        let x = 4.09;
        let y = 1.28;
        println!("x = {}, y = {}", x, y);
        x / y
    };
    println!("expr = {}", expr);

    // if-statement (every if-statement is an expression too, which is really a cool thing)
    let speed = 10;
    let mass = 890;
    let consumption =
        if speed < 0 {
            -1.0
        }
        else if speed == 0 {
            0.8
        }
        else if speed > 0 && speed < 10 {
            18.52
        }
        else if speed >= 10 && speed < 30 {
            let x = mass / speed;
            (x as f32) * 0.05
        }
        else {
            11.71
        };
        println!("Speed = {}, consumption = {}", speed, consumption);

        // loop
        let mut count = 0u32;
        loop {
            println!("loop iteration #{}", count);
            count += 1;
            if count > 9 {
            break;
            }
        }

        count = 0;
        while count < 10 {
            println!("while loop iteration #{}", count);
            count += 1;
        }

        for n in 0u32..10 {
            println!("for loop iteration #{}", n);
        }

        // matching
        let age = 19;
        match age {
            1  => println!("Seems like 1"),
            5  => println!("It's 5"),
            18 => println!("Exactly 18!"),
            _  => println!("I couldn't guess :("),
        }

        match age {
            a if a > 1 && a <= 9  => println!("The age is in range [1,9]"),
            a if a > 9 && a <= 19 => println!("The age is in range [10, 19]"),
            _                     => println!("Nothing, I don't know what range the age is in :("),
        }

        // the previous match is equivalent to the following:
        match age {
            a @ 1...9   => println!("The age is {} and in range [1,9]", a),
            a @ 10...19 => println!("The age is {} and in range [10, 19]", a),
            a           => println!("Nothing, I don't know what range the age {} is in :(", a),
        }

        // if let
        let a_some = Some(0.0008);
        if let Some(i) = a_some {
            println!("a_some is not None, {}", i);
            println!("result {}",  9.0 / i);
        }
        else {
            println!("None");
        }

        // while let
        let mut a_optional = Some(10);

        while let Some(i) = a_optional {
            if i == 20 {
                println!("while let: i equals 20, leave");
                a_optional = None;
            }
            else {
                a_optional = Some(i + 1);
            }
        }

        // functions, finally!
        let first_integer: i32 = -20;
        let second_integer: i32 = 8;
        let result_of_addition = add_two_integers(first_integer, second_integer);
        println!("{} + {} = {}", first_integer, second_integer, result_of_addition);

        simple_functions::print_hello_world();
        let mul_result = simple_functions::simple_math::multiple_two_integers(3, 15);
        println!("Modules: multiplication result = {}", mul_result);

        // this doesn't work.. why?
        println!("Modules: multiplication result 2 = {}", multiple_two_integers(-1, 8));

        // "super" and "self"
        sample_module::call_sample_module();

        // file hierarchy
        {
            let a : i32 = 4;
            let b : i32 = 15;
            println!("modules::math {} + {} = {}", a, b, modules::math::add_two_int(a, b));
        }

        // call external library
        my_extern_lib::run(5);

        // tuples
        let pair = ("key", 4);
        println!("tuple value: ({}, {})", pair.0, pair.1);
        println!("tuple value: {:?}", pair);

        // structs
        let bmw: Car = Car {color: "Blue".to_string(), weight: 940};
        println!("BMW: {} - {}", bmw.color, bmw.weight);

        // tuple structs
        let color_1: RgbColor = RgbColor(123, 33, 81);
        println!("RGB color: {}, {}, {}", color_1.0, color_1.1, color_1.2);

        // unit structs
        let _nil = Nil;

        // a struct with a constructor and getter
        let rabbit = my_extern_lib::Animal::new(2, 4, "The King");
        println!("Rabbit: {}, {}, {}", rabbit.age, rabbit.weight, rabbit.get_secret_info());

        // generics
        {
            struct ThreeValues<T, U, V> {
                first: T,
                second: U,
                third: V,
            }

            let three_values = ThreeValues { first: 4, second: 2.01, third: "hello" };
            println!("Generic data structure: {}, {}, {}", three_values.first, three_values.second, three_values.third);

            // generic function
            fn swap<U, V>(first_value: U, second_value: V) -> (V, U) {
                (second_value, first_value)
            }

            let random_tuple = swap(4, 2);
            println!("Use unused variables: {} {}", random_tuple.0, random_tuple.1);

        }

        // box (heap allocation)
        {

            let toyota: Car = Car {color: "Green".to_string(), weight: 1040,};
            let porche: Box<Car> = Box::new(Car {color: "Black".to_string(), weight: 876, });

            println!("Toyota's size on stack is {}", std::mem::size_of_val(&toyota));
            println!("Boxed porche's size on stack is {}", std::mem::size_of_val(&porche));

            let unboxed_porche: Car = *porche;
            println!("Unboxed porche's size on stack is {}", std::mem::size_of_val(&unboxed_porche));

            // ownership
            let bmw : Box<Car> = Box::new(Car {color: "Yellow".to_string(), weight: 1620});
            println!("Boxed BMW at {:p}: {}/{}", &bmw, bmw.color, bmw.weight);
            let moved_bmw: Box<Car> = bmw;
            println!("Moved BMW at {:p}: {}/{}", &moved_bmw, moved_bmw.color, moved_bmw.weight);
            // can't do the following, the ownership has been moved
            // println!("Boxed BMW at {:p}: {}/{}", &bmw, bmw.color, bmw.weight);

            // mutability
            let immutable_car = Box::new(Car {color: "Pink".to_string(), weight: 510});
            // can't do this
            // immutable_car.weight = 403;

            // but we can change mutability
            let mut mutable_car = immutable_car;
            mutable_car.weight = 401;

        }

        // ref
        {

            let mut bmw: Car = Car {color: "Green".to_string(), weight: 870};
            println!("Ref: bmw {},{}", bmw.color, bmw.weight);

            {
                let Car {color: _, weight: ref weight_reference} = bmw;

                println!("Ref: {}", *weight_reference);
            }

            {
                let Car {color: _, weight: ref mut weight_mut_ref} = bmw;

                *weight_mut_ref = 200;
            }

            println!("Ref: bmw {},{}", bmw.color, bmw.weight);

        }

        // try some enums
        {

            enum TwoValues {
                First(i8),
                Second(f32),
            }

            fn test_enum(some_number: i32) -> TwoValues {

                if some_number > 0 {
                    TwoValues::First(120)
                }
                else {
                    TwoValues::Second(3.09)
                }

            }

            match test_enum(-2) {
                TwoValues::First(value) => println!("TwoValues enum: {}", value),
                TwoValues::Second(value) => println!("TwoValues enum: {}", value),
            }

            let two_values_instance = TwoValues::First(9);
            println!("TwoValues enum: size of a variant {}", std::mem::size_of_val(&two_values_instance));

        }

        // closures
        {
            let say = |phrase: String| {println!("{}", phrase)};
            say("Hello, closure!".to_string()); // this is probably worse than just println!

            fn mutate_integer<F: Fn(i32) -> i32>(value: i32, mutator: F) -> i32 {
                mutator(value)
            }

            let modulo_three = |v: i32| {v % 3};

            let value: i32 = 23;
            println!("Closures: input value = {}, result = {}", value, mutate_integer(value, modulo_three));

            let add_five = |v: u32| {v + 5};
            let mul_ten = |v: u32| {v * 10};

            fn mutate_integer_2<F, G>(value: u32, mutator_1: F, mutator_2: G) -> u32
                where F: Fn(u32) -> u32, G: Fn(u32) -> u32 {

                mutator_2(mutator_1(value))

            }

            let value_2: u32 = 7194682;
            println!("Closures #2: input value = {}, result = {}", value_2, mutate_integer_2(value_2, add_five, mul_ten));
        }

        // threads
        {
            use std::thread;
            use std::old_io::timer;
            use std::sync::mpsc;
            use std::old_io;
            use std::time::Duration;

            // without a guard
            thread::spawn(|| {
                println!("Very first thread EEEEEEYYYYYEEEECCCCAAAAAAATTTCCCHHHEEEERRR!!!");
            });
            timer::sleep(Duration::milliseconds(50)); // we need it because main terminates and kills all the threads

            // with a guard
            let thread_guard = thread::scoped(|| {
                println!("Thread with a guard here!!");
            });

            // another example
            for i in range(0, 16) {

                thread::spawn(move || {
                    println!("Thread #{} started", i);
                    timer::sleep(Duration::milliseconds(500 % (i + 1)));
                    println!("Thread #{} finished", i);
                });

            }
            timer::sleep(Duration::milliseconds(2500));

            // channels

            let (tx, rx) = mpsc::channel();

            thread::spawn(move || {
                for idx in 10..20 {
                    tx.send(old_io::stdin().read_line().ok().expect("failed to read line"));
                }
            });

            for _ in 0..5 {
                println!("Got value {:?} from another thread", rx.recv().unwrap());
            }

        }
}

// a C-like struct
struct Car {
    color: String,
    weight: u32,
}

// a tuple struct s
struct RgbColor(i8, i8, i8);

// a unit struct
struct Nil;

fn add_two_integers(a: i32, b: i32) -> i32 {
    a + b
}
<|endoftext|>


<commit_before>use distribution::{Discrete, Distribution};
use random::Source;

/// A categorical distribution.
#[derive(Clone)]
pub struct Categorical {
    k: usize,
    p: Vec<f64>,
}

impl Categorical {
    /// Create a categorical distribution with success probability `p`.
    ///
    /// It should hold that `p[i] >= 0`, `p[i] <= 1`, and `sum(p) == 1`.
    #[inline]
    pub fn new(p: &[f64]) -> Categorical {
        should!(is_probability_vector(p), {
            const EPSILON: f64 = 1e-12;
            let mut in_unit = true;
            let mut sum = 0.0;
            for &p in p.iter() {
                if p < 0.0 || p > 1.0 {
                    in_unit = false;
                    break;
                }
                sum += p;
            }
            in_unit && (sum - 1.0).abs() <= EPSILON
        });
        Categorical { k: p.len(), p: p.to_vec() }
    }

    /// Return the number of categories.
    #[inline(always)]
    pub fn k(&self) -> usize { self.k }

    /// Return the event probabilities.
    #[inline(always)]
    pub fn p(&self) -> &[f64] { &self.p }
}

impl Distribution for Categorical {
    type Value = usize;

    fn mean(&self) -> f64 {
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| sum + i as f64 * p)
    }

    fn var(&self) -> f64 {
        let mean = self.mean();
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| sum + (i as f64 - mean).powi(2) * p)
    }

    fn skewness(&self) -> f64 {
        let (mean, var) = (self.mean(), self.var());
        let skew = self.p.iter().enumerate()
                                .fold(0.0, |sum, (i, p)| sum + (i as f64 - mean).powi(3) * p);
        skew / (var * var.sqrt())
    }

    fn kurtosis(&self) -> f64 {
        let (mean, var) = (self.mean(), self.var());
        let kurt = self.p.iter().enumerate()
                                .fold(0.0, |sum, (i, p)| sum + (i as f64 - mean).powi(4) * p);
        kurt / var.powi(2) - 3.0
    }

    fn median(&self) -> f64 {
        if self.p[0] > 0.5 {
            return 0.0;
        } else if self.p[0] == 0.5 {
            return 0.5;
        }
        let mut sum = 0.0;
        for i in 0..self.k {
            sum += self.p[i];
            if sum == 0.5 {
                return (2 * i - 1) as f64 / 2.0;
            } else if sum > 0.5 {
                return i as f64;
            }
        }
        unreachable!()
    }

    fn modes(&self) -> Vec<usize> {
        let mut modes = Vec::new();
        let mut max = 0.0;
        for (i, &p) in self.p.iter().enumerate() {
            if p == max {
                modes.push(i);
            }
            if p > max {
                max = p;
                modes = vec![i];
            }
        }
        modes
    }

    #[inline]
    fn entropy(&self) -> f64 {
        -self.p.iter().fold(0.0, |sum, p| sum + p * p.ln())
    }

    #[inline]
    fn cdf(&self, x: f64) -> f64 {
        if x < 0.0 {
            return 0.0;
        }
        let x = x as usize;
        if x >= self.k - 1 {
            1.0
        } else {
            self.p.iter().take(x + 1).fold(0.0, |a, b| a + b)
        }
    }

    fn inv_cdf(&self, p: f64) -> usize {
        should!(0.0 <= p && p <= 1.0);
        if p == 0.0 {
            return self.p.iter().position(|&p| p > 0.0).unwrap();
        }
        let mut sum = 0.0;
        for i in 0..self.k {
            sum += self.p[i];
            if sum >= p || sum == 1.0 {
                return i;
            }
        }
        self.k - 1
    }

    #[inline]
    fn pmf(&self, x: usize) -> f64 {
        self.p[x]
    }

    #[inline]
    fn sample<S>(&self, source: &mut S) -> usize where S: Source {
        self.inv_cdf(source.read::<f64>())
    }
}

impl Discrete for Categorical {
}

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

    macro_rules! new(
        (equal $k:expr) => { Categorical::new(&[1.0 / $k as f64; $k]) };
        ($p:expr) => { Categorical::new(&$p); }
    );

    #[test]
    fn mean() {
        assert_eq!(new!(equal 3).mean(), 1.0);
        assert_eq!(new!([0.3, 0.3, 0.4]).mean(), 1.1);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).mean(), 1.5);
    }

    #[test]
    fn var() {
        assert_eq!(new!(equal 3).var(), 2.0 / 3.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).var(), 11.0 / 12.0);
    }

    #[test]
    fn skewness() {
        assert_eq!(new!(equal 6).skewness(), 0.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).skewness(), 0.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).skewness(), -0.6);
    }

    #[test]
    fn kurtosis() {
        assert_eq!(new!(equal 2).kurtosis(), -2.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).kurtosis(), -0.7999999999999998);
    }

    #[test]
    fn median() {
        assert_eq!(new!([0.6, 0.2, 0.2]).median(), 0.0);
        assert_eq!(new!(equal 2).median(), 0.5);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).median(), 2.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).median(), 0.5);
    }

    #[test]
    fn modes() {
        assert_eq!(new!([0.6, 0.2, 0.2]).modes(), vec![0]);
        assert_eq!(new!(equal 2).modes(), vec![0, 1]);
        assert_eq!(new!(equal 3).modes(), vec![0, 1, 2]);
        assert_eq!(new!([0.4, 0.2, 0.4]).modes(), vec![0, 2]);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).modes(), vec![1, 2]);
    }

    #[test]
    fn entropy() {
        use std::f64::consts::LN_2;
        assert_eq!(new!(equal 2).entropy(), LN_2);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).entropy(), 1.2798542258336676);
    }

    #[test]
    fn pmf() {
        let p = [0.0, 0.75, 0.25, 0.0];
        let d1 = new!(p);
        assert_eq!(&(0..4).map(|x| d1.pmf(x)).collect::<Vec<_>>(), &p.to_vec());

        let d2 = new!(equal 3);
        assert_eq!(&(0..3).map(|x| d2.pmf(x)).collect::<Vec<_>>(), &vec![1.0 / 3.0; 3])
    }

    #[test]
    fn cdf() {
        let d = new!([0.0, 0.75, 0.25, 0.0]);
        let p = vec![0.0, 0.0, 0.75, 1.0, 1.0];

        let x = (-1..4).map(|x| d.cdf(x as f64)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let x = (-1..4).map(|x| d.cdf(x as f64 + 0.5)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let d = new!(equal 3);
        let p = vec![0.0, 1.0 / 3.0, 2.0 / 3.0, 1.0];

        let x = (-1..3).map(|x| d.cdf(x as f64)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let x = (-1..3).map(|x| d.cdf(x as f64 + 0.5)).collect::<Vec<_>>();
        assert_eq!(&x, &p);
    }

    #[test]
    fn inv_cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        let p1 = vec![0.0, 0.75, 0.7500001, 1.0];
        assert_eq!(&p1.iter().map(|&p| d1.inv_cdf(p)).collect::<Vec<_>>(), &vec![1, 1, 2, 2]);

        let d2 = new!(equal 3);
        let p2 = vec![0.0, 0.5, 0.75, 1.0];
        assert_eq!(&p2.iter().map(|&p| d2.inv_cdf(p)).collect::<Vec<_>>(), &vec![0, 1, 2, 2]);

    }

    #[test]
    fn sample() {
        let mut source = random::default();
        let sum = Independent(&new!([0.0, 0.5, 0.5]), &mut source).take(100).fold(0, |a, b| a + b);
        assert!(100 <= sum && sum <= 200);
    }
}
<commit_msg>categorical: refactor inv_cdf<commit_after>use distribution::{Discrete, Distribution};
use random::Source;

/// A categorical distribution.
#[derive(Clone)]
pub struct Categorical {
    k: usize,
    p: Vec<f64>,
}

impl Categorical {
    /// Create a categorical distribution with success probability `p`.
    ///
    /// It should hold that `p[i] >= 0`, `p[i] <= 1`, and `sum(p) == 1`.
    #[inline]
    pub fn new(p: &[f64]) -> Categorical {
        should!(is_probability_vector(p), {
            const EPSILON: f64 = 1e-12;
            let mut in_unit = true;
            let mut sum = 0.0;
            for &p in p.iter() {
                if p < 0.0 || p > 1.0 {
                    in_unit = false;
                    break;
                }
                sum += p;
            }
            in_unit && (sum - 1.0).abs() <= EPSILON
        });
        Categorical { k: p.len(), p: p.to_vec() }
    }

    /// Return the number of categories.
    #[inline(always)]
    pub fn k(&self) -> usize { self.k }

    /// Return the event probabilities.
    #[inline(always)]
    pub fn p(&self) -> &[f64] { &self.p }
}

impl Distribution for Categorical {
    type Value = usize;

    fn mean(&self) -> f64 {
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| sum + i as f64 * p)
    }

    fn var(&self) -> f64 {
        let mean = self.mean();
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| sum + (i as f64 - mean).powi(2) * p)
    }

    fn skewness(&self) -> f64 {
        let (mean, var) = (self.mean(), self.var());
        let skew = self.p.iter().enumerate()
                                .fold(0.0, |sum, (i, p)| sum + (i as f64 - mean).powi(3) * p);
        skew / (var * var.sqrt())
    }

    fn kurtosis(&self) -> f64 {
        let (mean, var) = (self.mean(), self.var());
        let kurt = self.p.iter().enumerate()
                                .fold(0.0, |sum, (i, p)| sum + (i as f64 - mean).powi(4) * p);
        kurt / var.powi(2) - 3.0
    }

    fn median(&self) -> f64 {
        if self.p[0] > 0.5 {
            return 0.0;
        } else if self.p[0] == 0.5 {
            return 0.5;
        }
        let mut sum = 0.0;
        for i in 0..self.k {
            sum += self.p[i];
            if sum == 0.5 {
                return (2 * i - 1) as f64 / 2.0;
            } else if sum > 0.5 {
                return i as f64;
            }
        }
        unreachable!()
    }

    fn modes(&self) -> Vec<usize> {
        let mut modes = Vec::new();
        let mut max = 0.0;
        for (i, &p) in self.p.iter().enumerate() {
            if p == max {
                modes.push(i);
            }
            if p > max {
                max = p;
                modes = vec![i];
            }
        }
        modes
    }

    #[inline]
    fn entropy(&self) -> f64 {
        -self.p.iter().fold(0.0, |sum, p| sum + p * p.ln())
    }

    #[inline]
    fn cdf(&self, x: f64) -> f64 {
        if x < 0.0 {
            return 0.0;
        }
        let x = x as usize;
        if x >= self.k - 1 {
            1.0
        } else {
            self.p.iter().take(x + 1).fold(0.0, |a, b| a + b)
        }
    }

    fn inv_cdf(&self, p: f64) -> usize {
        should!(0.0 <= p && p <= 1.0);
        let mut sum = 0.0;
        self.p.iter().position(|&pi| {
            sum += pi;
            pi > 0.0 && sum >= p
        }).unwrap_or_else(|| {
            self.p.iter().rposition(|&pi| pi > 0.0).unwrap()
        })
    }

    #[inline]
    fn pmf(&self, x: usize) -> f64 {
        self.p[x]
    }

    #[inline]
    fn sample<S>(&self, source: &mut S) -> usize where S: Source {
        self.inv_cdf(source.read::<f64>())
    }
}

impl Discrete for Categorical {
}

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

    macro_rules! new(
        (equal $k:expr) => { Categorical::new(&[1.0 / $k as f64; $k]) };
        ($p:expr) => { Categorical::new(&$p); }
    );

    #[test]
    fn mean() {
        assert_eq!(new!(equal 3).mean(), 1.0);
        assert_eq!(new!([0.3, 0.3, 0.4]).mean(), 1.1);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).mean(), 1.5);
    }

    #[test]
    fn var() {
        assert_eq!(new!(equal 3).var(), 2.0 / 3.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).var(), 11.0 / 12.0);
    }

    #[test]
    fn skewness() {
        assert_eq!(new!(equal 6).skewness(), 0.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).skewness(), 0.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).skewness(), -0.6);
    }

    #[test]
    fn kurtosis() {
        assert_eq!(new!(equal 2).kurtosis(), -2.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).kurtosis(), -0.7999999999999998);
    }

    #[test]
    fn median() {
        assert_eq!(new!([0.6, 0.2, 0.2]).median(), 0.0);
        assert_eq!(new!(equal 2).median(), 0.5);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).median(), 2.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).median(), 0.5);
    }

    #[test]
    fn modes() {
        assert_eq!(new!([0.6, 0.2, 0.2]).modes(), vec![0]);
        assert_eq!(new!(equal 2).modes(), vec![0, 1]);
        assert_eq!(new!(equal 3).modes(), vec![0, 1, 2]);
        assert_eq!(new!([0.4, 0.2, 0.4]).modes(), vec![0, 2]);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).modes(), vec![1, 2]);
    }

    #[test]
    fn entropy() {
        use std::f64::consts::LN_2;
        assert_eq!(new!(equal 2).entropy(), LN_2);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).entropy(), 1.2798542258336676);
    }

    #[test]
    fn pmf() {
        let p = [0.0, 0.75, 0.25, 0.0];
        let d1 = new!(p);
        assert_eq!(&(0..4).map(|x| d1.pmf(x)).collect::<Vec<_>>(), &p.to_vec());

        let d2 = new!(equal 3);
        assert_eq!(&(0..3).map(|x| d2.pmf(x)).collect::<Vec<_>>(), &vec![1.0 / 3.0; 3])
    }

    #[test]
    fn cdf() {
        let d = new!([0.0, 0.75, 0.25, 0.0]);
        let p = vec![0.0, 0.0, 0.75, 1.0, 1.0];

        let x = (-1..4).map(|x| d.cdf(x as f64)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let x = (-1..4).map(|x| d.cdf(x as f64 + 0.5)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let d = new!(equal 3);
        let p = vec![0.0, 1.0 / 3.0, 2.0 / 3.0, 1.0];

        let x = (-1..3).map(|x| d.cdf(x as f64)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let x = (-1..3).map(|x| d.cdf(x as f64 + 0.5)).collect::<Vec<_>>();
        assert_eq!(&x, &p);
    }

    #[test]
    fn inv_cdf() {
        let d1 = new!([0.0, 0.75, 0.25, 0.0]);
        let p1 = vec![0.0, 0.75, 0.7500001, 1.0];
        assert_eq!(&p1.iter().map(|&p| d1.inv_cdf(p)).collect::<Vec<_>>(), &vec![1, 1, 2, 2]);

        let d2 = new!(equal 3);
        let p2 = vec![0.0, 0.5, 0.75, 1.0];
        assert_eq!(&p2.iter().map(|&p| d2.inv_cdf(p)).collect::<Vec<_>>(), &vec![0, 1, 2, 2]);

    }

    #[test]
    fn sample() {
        let mut source = random::default();
        let sum = Independent(&new!([0.0, 0.5, 0.5]), &mut source).take(100).fold(0, |a, b| a + b);
        assert!(100 <= sum && sum <= 200);
    }
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use devicemapper::DM;

use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::iter::FromIterator;
use std::path::Path;
use std::path::PathBuf;
use std::vec::Vec;

use time;
use uuid::Uuid;
use serde_json;

use engine::EngineError;
use engine::EngineResult;
use engine::ErrorEnum;
use engine::Filesystem;
use engine::Pool;
use engine::RenameAction;
use engine::engine::Redundancy;
use engine::strat_engine::lineardev::LinearDev;
use engine::strat_engine::thinpooldev::ThinPoolDev;

use super::super::engine::{FilesystemUuid, HasName, HasUuid};
use super::super::structures::Table;

use super::serde_structs::StratSave;
use super::blockdev::{BlockDev, initialize, resolve_devices};
use super::filesystem::StratFilesystem;
use super::metadata::MIN_MDA_SECTORS;

use types::DataBlocks;
use types::Sectors;

#[derive(Debug)]
pub struct StratPool {
    name: String,
    pool_uuid: Uuid,
    pub block_devs: BTreeMap<PathBuf, BlockDev>,
    pub filesystems: Table<StratFilesystem>,
    redundancy: Redundancy,
    thin_pool: ThinPoolDev,
}

impl StratPool {
    pub fn new(name: &str,
               dm: &DM,
               paths: &[&Path],
               redundancy: Redundancy,
               force: bool)
               -> EngineResult<StratPool> {
        let pool_uuid = Uuid::new_v4();

        let devices = try!(resolve_devices(paths));
        let bds = try!(initialize(&pool_uuid, devices, MIN_MDA_SECTORS, force));

        // TODO: We've got some temporary code in BlockDev::initialize that
        // makes sure we've got at least 2 blockdevs supplied - one for a meta
        // and one for data.  In the future, we will be able to deal with a
        // single blockdev.  When that code is added to use a single blockdev,
        // the check for 2 devs in BlockDev::initialize should be removed.
        assert!(bds.len() >= 2);
        let meta_dev;
        let data_dev;
        let length;
        {
            let mut data_devs = Vec::from_iter(bds.values());

            meta_dev = try!(LinearDev::new(&format!("stratis_{}_meta", name),
                                           dm,
                                           &vec![data_devs.remove(0)]));

            data_dev = try!(LinearDev::new(&format!("stratis_{}_data", name), dm, &data_devs));
            length = try!(data_dev.size()).sectors();

        }
        // TODO Fix hard coded data blocksize and low water mark.
        let thinpool_dev = try!(ThinPoolDev::new(&format!("stratis_{}_thinpool", name),
                                                 dm,
                                                 length,
                                                 Sectors(1024),
                                                 DataBlocks(256000),
                                                 meta_dev,
                                                 data_dev));

        let mut pool = StratPool {
            name: name.to_owned(),
            pool_uuid: pool_uuid,
            block_devs: bds,
            filesystems: Table::new(),
            redundancy: redundancy,
            thin_pool: thinpool_dev,
        };

        try!(pool.write_metadata());

        Ok(pool)
    }

    /// Return the metadata from the first blockdev with up-to-date, readable
    /// metadata.
    pub fn read_metadata(&self) -> Option<Vec<u8>> {

        let mut bds: Vec<&BlockDev> = self.block_devs
            .iter()
            .map(|(_, bd)| bd)
            .filter(|bd| bd.last_update_time().is_some())
            .collect();

        if bds.is_empty() {
            return None;
        }

        bds.sort_by_key(|k| {
            k.last_update_time().expect("devs without some last update time filtered above.")
        });

        let last_update_time = bds.last()
            .expect("There is a last bd, since bds.is_empty() was false.")
            .last_update_time()
            .expect("devs without some last update time filtered above.");

        for bd in bds.iter()
            .rev()
            .take_while(|bd| {
                bd.last_update_time()
                    .expect("devs without some last update time filtered above.") ==
                last_update_time
            }) {
            match bd.load_state() {
                Ok(Some(data)) => return Some(data),
                _ => continue,
            }
        }
        None
    }

    /// Write the given data to all blockdevs
    pub fn write_metadata(&mut self) -> EngineResult<()> {

        let data = try!(serde_json::to_string(&self.to_save()));

        // TODO: Cap # of blockdevs written to, as described in SWDD
        // TODO: Check current time against global last updated, and use
        // alternate time value if earlier, as described in SWDD

        let time = time::now().to_timespec();

        // TODO: Do something better than panic when saving to blockdev fails.
        for (_, bd) in &mut self.block_devs {
            bd.save_state(&time, data.as_bytes()).unwrap();
        }

        Ok(())
    }

    pub fn to_save(&self) -> StratSave {
        StratSave {
            name: self.name.clone(),
            id: self.pool_uuid.simple().to_string(),
            block_devs: self.block_devs
                .iter()
                .map(|(_, bd)| (bd.uuid().simple().to_string(), bd.to_save()))
                .collect(),
        }
    }
}

impl Pool for StratPool {
    fn create_filesystems<'a, 'b>(&'a mut self,
                                  specs: &[&'b str])
                                  -> EngineResult<Vec<(&'b str, FilesystemUuid)>> {
        let names = BTreeSet::from_iter(specs);
        for name in names.iter() {
            if self.filesystems.contains_name(name) {
                return Err(EngineError::Engine(ErrorEnum::AlreadyExists, name.to_string()));
            }
        }

        let mut result = Vec::new();
        for name in names.iter() {
            let uuid = Uuid::new_v4();
            let dm = try!(DM::new());
            let new_filesystem = try!(StratFilesystem::new(uuid, name, &dm, &mut self.thin_pool));
            self.filesystems.insert(new_filesystem);
            result.push((**name, uuid));
        }

        Ok(result)
    }

    fn add_blockdevs(&mut self, paths: &[&Path], force: bool) -> EngineResult<Vec<PathBuf>> {
        let devices = try!(resolve_devices(paths));
        let mut bds = try!(initialize(&self.pool_uuid, devices, MIN_MDA_SECTORS, force));
        let bdev_paths = bds.iter().map(|p| p.1.devnode.clone()).collect();
        self.block_devs.append(&mut bds);
        try!(self.write_metadata());
        Ok(bdev_paths)
    }

    fn destroy(mut self) -> EngineResult<()> {

        // TODO Do we want to create a new File each time we interact with DM?
        let dm = try!(DM::new());
        try!(self.thin_pool.teardown(&dm));

        for bd in self.block_devs.values_mut() {
            try!(bd.wipe_metadata());
        }

        Ok(())
    }

    fn destroy_filesystems<'a, 'b>(&'a mut self,
                                   fs_uuids: &[&'b FilesystemUuid])
                                   -> EngineResult<Vec<&'b FilesystemUuid>> {
        destroy_filesystems!{self; fs_uuids}
    }

    fn rename_filesystem(&mut self,
                         uuid: &FilesystemUuid,
                         new_name: &str)
                         -> EngineResult<RenameAction> {
        rename_filesystem!{self; uuid; new_name}
    }

    fn rename(&mut self, name: &str) {
        self.name = name.to_owned();
    }

    fn get_filesystem(&mut self, uuid: &FilesystemUuid) -> Option<&mut Filesystem> {
        get_filesystem!(self; uuid)
    }
}

impl HasUuid for StratPool {
    fn uuid(&self) -> &FilesystemUuid {
        &self.pool_uuid
    }
}

impl HasName for StratPool {
    fn name(&self) -> &str {
        &self.name
    }
}
<commit_msg>Change signature of StratPool::read_metadata<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use devicemapper::DM;

use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::iter::FromIterator;
use std::path::Path;
use std::path::PathBuf;
use std::vec::Vec;

use time;
use uuid::Uuid;
use serde_json;

use engine::EngineError;
use engine::EngineResult;
use engine::ErrorEnum;
use engine::Filesystem;
use engine::Pool;
use engine::RenameAction;
use engine::engine::Redundancy;
use engine::strat_engine::lineardev::LinearDev;
use engine::strat_engine::thinpooldev::ThinPoolDev;

use super::super::engine::{FilesystemUuid, HasName, HasUuid};
use super::super::structures::Table;

use super::serde_structs::StratSave;
use super::blockdev::{BlockDev, initialize, resolve_devices};
use super::filesystem::StratFilesystem;
use super::metadata::MIN_MDA_SECTORS;

use types::DataBlocks;
use types::Sectors;

#[derive(Debug)]
pub struct StratPool {
    name: String,
    pool_uuid: Uuid,
    pub block_devs: BTreeMap<PathBuf, BlockDev>,
    pub filesystems: Table<StratFilesystem>,
    redundancy: Redundancy,
    thin_pool: ThinPoolDev,
}

impl StratPool {
    pub fn new(name: &str,
               dm: &DM,
               paths: &[&Path],
               redundancy: Redundancy,
               force: bool)
               -> EngineResult<StratPool> {
        let pool_uuid = Uuid::new_v4();

        let devices = try!(resolve_devices(paths));
        let bds = try!(initialize(&pool_uuid, devices, MIN_MDA_SECTORS, force));

        // TODO: We've got some temporary code in BlockDev::initialize that
        // makes sure we've got at least 2 blockdevs supplied - one for a meta
        // and one for data.  In the future, we will be able to deal with a
        // single blockdev.  When that code is added to use a single blockdev,
        // the check for 2 devs in BlockDev::initialize should be removed.
        assert!(bds.len() >= 2);
        let meta_dev;
        let data_dev;
        let length;
        {
            let mut data_devs = Vec::from_iter(bds.values());

            meta_dev = try!(LinearDev::new(&format!("stratis_{}_meta", name),
                                           dm,
                                           &vec![data_devs.remove(0)]));

            data_dev = try!(LinearDev::new(&format!("stratis_{}_data", name), dm, &data_devs));
            length = try!(data_dev.size()).sectors();

        }
        // TODO Fix hard coded data blocksize and low water mark.
        let thinpool_dev = try!(ThinPoolDev::new(&format!("stratis_{}_thinpool", name),
                                                 dm,
                                                 length,
                                                 Sectors(1024),
                                                 DataBlocks(256000),
                                                 meta_dev,
                                                 data_dev));

        let mut pool = StratPool {
            name: name.to_owned(),
            pool_uuid: pool_uuid,
            block_devs: bds,
            filesystems: Table::new(),
            redundancy: redundancy,
            thin_pool: thinpool_dev,
        };

        try!(pool.write_metadata());

        Ok(pool)
    }

    /// Return the metadata from the first blockdev with up-to-date, readable
    /// metadata.
    /// Precondition: All Blockdevs in blockdevs must belong to the same pool.
    pub fn read_metadata(blockdevs: &[&BlockDev]) -> Option<Vec<u8>> {

        let mut bds: Vec<_> = blockdevs.iter()
            .filter(|bd| bd.last_update_time().is_some())
            .collect();

        if bds.is_empty() {
            return None;
        }

        bds.sort_by_key(|k| {
            k.last_update_time().expect("devs without some last update time filtered above.")
        });

        let last_update_time = bds.last()
            .expect("There is a last bd, since bds.is_empty() was false.")
            .last_update_time()
            .expect("devs without some last update time filtered above.");

        for bd in bds.iter()
            .rev()
            .take_while(|bd| {
                bd.last_update_time()
                    .expect("devs without some last update time filtered above.") ==
                last_update_time
            }) {
            match bd.load_state() {
                Ok(Some(data)) => return Some(data),
                _ => continue,
            }
        }
        None
    }

    /// Write the given data to all blockdevs
    pub fn write_metadata(&mut self) -> EngineResult<()> {

        let data = try!(serde_json::to_string(&self.to_save()));

        // TODO: Cap # of blockdevs written to, as described in SWDD
        // TODO: Check current time against global last updated, and use
        // alternate time value if earlier, as described in SWDD

        let time = time::now().to_timespec();

        // TODO: Do something better than panic when saving to blockdev fails.
        for (_, bd) in &mut self.block_devs {
            bd.save_state(&time, data.as_bytes()).unwrap();
        }

        Ok(())
    }

    pub fn to_save(&self) -> StratSave {
        StratSave {
            name: self.name.clone(),
            id: self.pool_uuid.simple().to_string(),
            block_devs: self.block_devs
                .iter()
                .map(|(_, bd)| (bd.uuid().simple().to_string(), bd.to_save()))
                .collect(),
        }
    }
}

impl Pool for StratPool {
    fn create_filesystems<'a, 'b>(&'a mut self,
                                  specs: &[&'b str])
                                  -> EngineResult<Vec<(&'b str, FilesystemUuid)>> {
        let names = BTreeSet::from_iter(specs);
        for name in names.iter() {
            if self.filesystems.contains_name(name) {
                return Err(EngineError::Engine(ErrorEnum::AlreadyExists, name.to_string()));
            }
        }

        let mut result = Vec::new();
        for name in names.iter() {
            let uuid = Uuid::new_v4();
            let dm = try!(DM::new());
            let new_filesystem = try!(StratFilesystem::new(uuid, name, &dm, &mut self.thin_pool));
            self.filesystems.insert(new_filesystem);
            result.push((**name, uuid));
        }

        Ok(result)
    }

    fn add_blockdevs(&mut self, paths: &[&Path], force: bool) -> EngineResult<Vec<PathBuf>> {
        let devices = try!(resolve_devices(paths));
        let mut bds = try!(initialize(&self.pool_uuid, devices, MIN_MDA_SECTORS, force));
        let bdev_paths = bds.iter().map(|p| p.1.devnode.clone()).collect();
        self.block_devs.append(&mut bds);
        try!(self.write_metadata());
        Ok(bdev_paths)
    }

    fn destroy(mut self) -> EngineResult<()> {

        // TODO Do we want to create a new File each time we interact with DM?
        let dm = try!(DM::new());
        try!(self.thin_pool.teardown(&dm));

        for bd in self.block_devs.values_mut() {
            try!(bd.wipe_metadata());
        }

        Ok(())
    }

    fn destroy_filesystems<'a, 'b>(&'a mut self,
                                   fs_uuids: &[&'b FilesystemUuid])
                                   -> EngineResult<Vec<&'b FilesystemUuid>> {
        destroy_filesystems!{self; fs_uuids}
    }

    fn rename_filesystem(&mut self,
                         uuid: &FilesystemUuid,
                         new_name: &str)
                         -> EngineResult<RenameAction> {
        rename_filesystem!{self; uuid; new_name}
    }

    fn rename(&mut self, name: &str) {
        self.name = name.to_owned();
    }

    fn get_filesystem(&mut self, uuid: &FilesystemUuid) -> Option<&mut Filesystem> {
        get_filesystem!(self; uuid)
    }
}

impl HasUuid for StratPool {
    fn uuid(&self) -> &FilesystemUuid {
        &self.pool_uuid
    }
}

impl HasName for StratPool {
    fn name(&self) -> &str {
        &self.name
    }
}
<|endoftext|>


<commit_before>//! Geometric queries.

#![deny(non_camel_case_types)]
#![deny(unused_parens)]
#![deny(non_upper_case_globals)]
#![deny(unused_qualifications)]
#![warn(missing_docs)]
#![deny(unused_results)]
#![warn(unused_imports)]
#![deny(unused_typecasts)]
#![feature(unsafe_destructor)]
#![feature(plugin)]
#![feature(old_impl_check)]
#![feature(old_orphan_check)]
#![doc(html_root_url = "http://ncollide.org/doc")]

extern crate test; // To compute the median.
extern crate "rustc-serialize" as rustc_serialize;
extern crate "nalgebra" as na;
extern crate "ncollide_math" as math;
extern crate "ncollide_utils" as utils;
#[plugin]
extern crate ncollide_entities;
extern crate "ncollide_entities" as entities;

pub mod geometry;
pub mod point;
pub mod ray;
<commit_msg>Update plugin syntax.<commit_after>//! Geometric queries.

#![deny(non_camel_case_types)]
#![deny(unused_parens)]
#![deny(non_upper_case_globals)]
#![deny(unused_qualifications)]
#![warn(missing_docs)]
#![deny(unused_results)]
#![warn(unused_imports)]
#![deny(unused_typecasts)]
#![feature(unsafe_destructor)]
#![feature(plugin)]
#![feature(old_impl_check)]
#![feature(old_orphan_check)]
#![doc(html_root_url = "http://ncollide.org/doc")]

#![plugin(ncollide_entities)]

extern crate test; // To compute the median.
extern crate "rustc-serialize" as rustc_serialize;
extern crate "nalgebra" as na;
extern crate "ncollide_math" as math;
extern crate "ncollide_utils" as utils;
extern crate ncollide_entities;
extern crate "ncollide_entities" as entities;

pub mod geometry;
pub mod point;
pub mod ray;
<|endoftext|>


<commit_before>use crate::error::Error;
use std::collections::HashMap;

pub fn encoded_len(input_len: usize) -> usize {
    if input_len == 0 {
        0
    } else {
        (input_len * 8 - 1) / 5 + 1
    }
}

static BASE32_CHARS: &'static [u8; 32] = &b"0123456789abcdfghijklmnpqrsvwxyz";

lazy_static! {
    static ref BASE32_CHARS_REVERSE: HashMap<char, u8> = BASE32_CHARS
        .iter()
        .enumerate()
        .map(|(n, c)| (*c as char, n as u8))
        .collect();
}

pub fn encode(input: &[u8]) -> String {
    let mut res = String::new();
    res.reserve(encoded_len(input.len()));

    let mut nr_bits_left: usize = 0;
    let mut bits_left: u16 = 0;

    for b in input {
        bits_left |= (*b as u16) << nr_bits_left;
        nr_bits_left += 8;
        while nr_bits_left > 5 {
            res.push(BASE32_CHARS[(bits_left & 0x1f) as usize] as char);
            bits_left >>= 5;
            nr_bits_left -= 5;
        }
    }

    if nr_bits_left > 0 {
        res.push(BASE32_CHARS[(bits_left & 0x1f) as usize] as char);
    }

    res.chars().rev().collect()
}

pub fn decode(input: &str) -> Result<Vec<u8>, crate::Error> {
    let mut res = Vec::new();

    let mut nr_bits_left: usize = 0;
    let mut bits_left: u16 = 0;

    for c in input.chars().rev() {
        let x = BASE32_CHARS_REVERSE.get(&c).ok_or(Error::BadBase32)?;
        bits_left |= (*x as u16) << nr_bits_left;
        nr_bits_left += 5;
        if nr_bits_left >= 8 {
            res.push((bits_left & 0xff) as u8);
            bits_left >>= 8;
            nr_bits_left -= 8;
        }
    }

    if nr_bits_left > 0 && bits_left != 0 {
        return Err(Error::BadBase32);
    }

    Ok(res)
}

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

    #[test]
    fn test_encode() {
        assert_eq!(encode(&[]), "");

        assert_eq!(
            encode(&hex::decode("0839703786356bca59b0f4a32987eb2e6de43ae8").unwrap()),
            "x0xf8v9fxf3jk8zln1cwlsrmhqvp0f88"
        );

        assert_eq!(
            encode(
                &hex::decode("ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
                    .unwrap()
            ),
            "1b8m03r63zqhnjf7l5wnldhh7c134ap5vpj0850ymkq1iyzicy5s"
        );

        assert_eq!(
            encode(
                &hex::decode("ddaf35a193617abacc417349ae20413112e6fa4e89a97ea20a9eeee64b55d39a2192992a274fc1a836ba3c23a3feebbd454d4423643ce80e2a9ac94fa54ca49f")
                    .unwrap()
            ),
            "2gs8k559z4rlahfx0y688s49m2vvszylcikrfinm30ly9rak69236nkam5ydvly1ai7xac99vxfc4ii84hawjbk876blyk1jfhkbbyx"
        );
    }

    #[test]
    fn test_decode() {
        assert_eq!(hex::encode(decode("").unwrap()), "");

        assert_eq!(
            hex::encode(decode("x0xf8v9fxf3jk8zln1cwlsrmhqvp0f88").unwrap()),
            "0839703786356bca59b0f4a32987eb2e6de43ae8"
        );

        assert_eq!(
            hex::encode(decode("1b8m03r63zqhnjf7l5wnldhh7c134ap5vpj0850ymkq1iyzicy5s").unwrap()),
            "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"
        );

        assert_eq!(
            hex::encode(decode("2gs8k559z4rlahfx0y688s49m2vvszylcikrfinm30ly9rak69236nkam5ydvly1ai7xac99vxfc4ii84hawjbk876blyk1jfhkbbyx").unwrap()),
            "ddaf35a193617abacc417349ae20413112e6fa4e89a97ea20a9eeee64b55d39a2192992a274fc1a836ba3c23a3feebbd454d4423643ce80e2a9ac94fa54ca49f"
        );

        assert_matches!(
            decode("xoxf8v9fxf3jk8zln1cwlsrmhqvp0f88"),
            Err(Error::BadBase32)
        );
        assert_matches!(
            decode("2b8m03r63zqhnjf7l5wnldhh7c134ap5vpj0850ymkq1iyzicy5s"),
            Err(Error::BadBase32)
        );
        assert_matches!(decode("2"), Err(Error::BadBase32));
        assert_matches!(decode("2gs"), Err(Error::BadBase32));
        assert_matches!(decode("2gs8"), Err(Error::BadBase32));
    }

    proptest! {

        #[test]
        fn roundtrip(s: Vec<u8>) {
            assert_eq!(s, decode(&encode(&s)).unwrap());
        }
    }
}
<commit_msg>Speed up base32::decode()<commit_after>use crate::error::Error;

pub fn encoded_len(input_len: usize) -> usize {
    if input_len == 0 {
        0
    } else {
        (input_len * 8 - 1) / 5 + 1
    }
}

pub fn decoded_len(input_len: usize) -> usize {
    input_len * 5 / 8
}

static BASE32_CHARS: &'static [u8; 32] = &b"0123456789abcdfghijklmnpqrsvwxyz";

lazy_static! {
    static ref BASE32_CHARS_REVERSE: Box<[u8; 256]> = {
        let mut xs = [0xffu8; 256];
        for (n, c) in BASE32_CHARS.iter().enumerate() {
            xs[*c as usize] = n as u8;
        }
        Box::new(xs)
    };
}

pub fn encode(input: &[u8]) -> String {
    let mut res = String::new();
    res.reserve(encoded_len(input.len()));

    let mut nr_bits_left: usize = 0;
    let mut bits_left: u16 = 0;

    for b in input {
        bits_left |= (*b as u16) << nr_bits_left;
        nr_bits_left += 8;
        while nr_bits_left > 5 {
            res.push(BASE32_CHARS[(bits_left & 0x1f) as usize] as char);
            bits_left >>= 5;
            nr_bits_left -= 5;
        }
    }

    if nr_bits_left > 0 {
        res.push(BASE32_CHARS[(bits_left & 0x1f) as usize] as char);
    }

    res.chars().rev().collect()
}

pub fn decode(input: &str) -> Result<Vec<u8>, crate::Error> {
    let mut res = Vec::with_capacity(decoded_len(input.len()));

    let mut nr_bits_left: usize = 0;
    let mut bits_left: u16 = 0;

    for c in input.chars().rev() {
        let b = BASE32_CHARS_REVERSE[c as usize];
        if b == 0xff {
            return Err(Error::BadBase32);
        }
        bits_left |= (b as u16) << nr_bits_left;
        nr_bits_left += 5;
        if nr_bits_left >= 8 {
            res.push((bits_left & 0xff) as u8);
            bits_left >>= 8;
            nr_bits_left -= 8;
        }
    }

    if nr_bits_left > 0 && bits_left != 0 {
        return Err(Error::BadBase32);
    }

    Ok(res)
}

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

    #[test]
    fn test_encode() {
        assert_eq!(encode(&[]), "");

        assert_eq!(
            encode(&hex::decode("0839703786356bca59b0f4a32987eb2e6de43ae8").unwrap()),
            "x0xf8v9fxf3jk8zln1cwlsrmhqvp0f88"
        );

        assert_eq!(
            encode(
                &hex::decode("ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")
                    .unwrap()
            ),
            "1b8m03r63zqhnjf7l5wnldhh7c134ap5vpj0850ymkq1iyzicy5s"
        );

        assert_eq!(
            encode(
                &hex::decode("ddaf35a193617abacc417349ae20413112e6fa4e89a97ea20a9eeee64b55d39a2192992a274fc1a836ba3c23a3feebbd454d4423643ce80e2a9ac94fa54ca49f")
                    .unwrap()
            ),
            "2gs8k559z4rlahfx0y688s49m2vvszylcikrfinm30ly9rak69236nkam5ydvly1ai7xac99vxfc4ii84hawjbk876blyk1jfhkbbyx"
        );
    }

    #[test]
    fn test_decode() {
        assert_eq!(hex::encode(decode("").unwrap()), "");

        assert_eq!(
            hex::encode(decode("x0xf8v9fxf3jk8zln1cwlsrmhqvp0f88").unwrap()),
            "0839703786356bca59b0f4a32987eb2e6de43ae8"
        );

        assert_eq!(
            hex::encode(decode("1b8m03r63zqhnjf7l5wnldhh7c134ap5vpj0850ymkq1iyzicy5s").unwrap()),
            "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"
        );

        assert_eq!(
            hex::encode(decode("2gs8k559z4rlahfx0y688s49m2vvszylcikrfinm30ly9rak69236nkam5ydvly1ai7xac99vxfc4ii84hawjbk876blyk1jfhkbbyx").unwrap()),
            "ddaf35a193617abacc417349ae20413112e6fa4e89a97ea20a9eeee64b55d39a2192992a274fc1a836ba3c23a3feebbd454d4423643ce80e2a9ac94fa54ca49f"
        );

        assert_matches!(
            decode("xoxf8v9fxf3jk8zln1cwlsrmhqvp0f88"),
            Err(Error::BadBase32)
        );
        assert_matches!(
            decode("2b8m03r63zqhnjf7l5wnldhh7c134ap5vpj0850ymkq1iyzicy5s"),
            Err(Error::BadBase32)
        );
        assert_matches!(decode("2"), Err(Error::BadBase32));
        assert_matches!(decode("2gs"), Err(Error::BadBase32));
        assert_matches!(decode("2gs8"), Err(Error::BadBase32));
    }

    proptest! {

        #[test]
        fn roundtrip(s: Vec<u8>) {
            assert_eq!(s, decode(&encode(&s)).unwrap());
        }
    }
}
<|endoftext|>