<commit_before>// Hound -- A WAV encoding library in Rust
// Copyright (C) 2015 Ruud van Asseldonk
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3 of the License only.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

//! Hound, a WAV encoding library.
//!
//! TODO: Add some introductory text here.
//!
//! The following example renders a 440 Hz sine wave, and stores it as as a
//! mono wav file with a sample rate of 44.1 kHz and 16 bits per sample.
//!
//! ```
//! use std::fs;
//! use std::f32::consts::PI;
//! use std::i16;
//! use hound;
//!
//! let spec = hound::WavSpec {
//!     channels: 1,
//!     sample_rate: 44100,
//!     bits_per_sample: 16
//! };
//! // TODO: ensure that the type can be inferred.
//! let writer_res = hound::WavWriter::<fs::File>::create("sine.wav", spec);
//! let mut writer = writer_res.ok().unwrap();
//! for t in (0 .. 44100).map(|x| x as f32 / 44100.0) {
//!     let sample = (t * 440.0 * 2.0 * PI).sin();
//!     let amplitude: i16 = i16::MAX;
//!     writer.write_sample((sample * amplitude as f32) as i16).ok().unwrap();
//! }
//! writer.finalize().ok().unwrap();
//! ```

#![warn(missing_docs)]

use std::fs;
use std::io;
use std::io::Write;
use std::path;

/// Extends the functionality of `io::Write` with additional methods.
///
/// The methods may be used on any type that implements `io::Write`.
trait WriteExt: io::Write {
    /// Writes an unsigned 16-bit integer in little endian format.
    fn write_le_u16(&mut self, x: u16) -> io::Result<()>;

    /// Writes an unsigned 32-bit integer in little endian format.
    fn write_le_u32(&mut self, x: u32) -> io::Result<()>;
}

impl<W> WriteExt for W where W: io::Write {

    fn write_le_u16(&mut self, x: u16) -> io::Result<()> {
        let mut buf = [0u8; 2];
        buf[0] = (x & 0xff) as u8;
        buf[1] = (x >> 8) as u8;
        self.write_all(&buf)
    }

    fn write_le_u32(&mut self, x: u32) -> io::Result<()> {
        let mut buf = [0u8; 4];
        buf[0] = ((x >> 00) & 0xff) as u8;
        buf[1] = ((x >> 08) & 0xff) as u8;
        buf[2] = ((x >> 16) & 0xff) as u8;
        buf[3] = ((x >> 24) & 0xff) as u8;
        self.write_all(&buf)
    }
}

/// A type that can be used to represent audio samples.
pub trait Sample {
    /// Writes the audio sample to the WAVE data section.
    fn write<W: io::Write>(self, writer: &mut W, bits: u32) -> io::Result<()>;
}

impl Sample for u16 {
    fn write<W: io::Write>(self, writer: &mut W, bits: u32) -> io::Result<()> {
        writer.write_le_u16(self)
        // TODO: take bits into account
    }
}

impl Sample for i16 {
    fn write<W: io::Write>(self, writer: &mut W, bits: u32) -> io::Result<()> {
        writer.write_le_u16(self as u16)
        // TODO: take bits into account
    }
}

/// Generates a bitmask with `channels` ones in the least significant bits.
fn channel_mask(channels: u16) -> u32 {
    (0 .. channels).map(|c| 1 << c).fold(0, |a, c| a | c)
}

#[test]
fn verify_channel_mask()
{
    assert_eq!(channel_mask(0), 0);
    assert_eq!(channel_mask(1), 1);
    assert_eq!(channel_mask(2), 3);
    assert_eq!(channel_mask(3), 7);
    assert_eq!(channel_mask(4), 15);
}

/// Specifies properties of the audio data.
#[derive(Clone, Copy)]
pub struct WavSpec {
    /// The number of channels.
    pub channels: u16,

    /// The number of samples per second.
    ///
    /// A common value is 44100, this is 44.1 kHz which is used for CD audio.
    pub sample_rate: u32,

    /// The number of bits per sample.
    ///
    /// A common value is 16 bits per sample, which is used for CD audio.
    pub bits_per_sample: u32
}

/// A writer that accepts samples and writes the WAVE format.
///
/// TODO: add example.
///
/// After all samples have been written, the file must be finalized. This can
/// be done by calling `finalize`. If `finalize` is not called, the file will
/// be finalized upon drop. However, finalization involves IO that may fail,
/// and without calling `finalize`, such a failure cannot be observed.
pub struct WavWriter<W> where W: io::Write + io::Seek {
    /// Specifies properties of the audio data.
    spec: WavSpec,

    /// The (container) bytes per sample. This is the bit rate / 8 rounded up.
    bytes_per_sample: u32,

    /// Whether the header has been written already.
    wrote_header: bool,

    /// The writer that will be written to.
    writer: io::BufWriter<W>,

    /// The number of bytes written to the data section.
    ///
    /// This is an `u32` because WAVE cannot accomodate more data.
    data_bytes_written: u32,

    /// Whether `finalize_internal` has been called.
    finalized: bool
}

impl<W> WavWriter<W> where W: io::Write + io::Seek {
    /// Creates a writer that writes the WAVE format to the underlying writer.
    ///
    /// The underlying writer is assumed to be at offset 0. `WavWriter` employs
    /// buffering internally to avoid too many `write` calls to the underlying
    /// writer.
    pub fn new(writer: W, spec: WavSpec) -> WavWriter<W> {
        WavWriter {
            spec: spec,
            bytes_per_sample: (spec.bits_per_sample as f32 / 8.0).ceil() as u32,
            wrote_header: false,
            writer: io::BufWriter::new(writer),
            data_bytes_written: 0,
            finalized: false
        }
    }

    /// Creates a writer that writes the WAVE format to a file.
    ///
    /// The file will be overwritten if it exists.
    pub fn create<P: AsRef<path::Path>>(filename: P, spec: WavSpec)
           -> io::Result<WavWriter<fs::File>> {
        let file = try!(fs::File::create(filename));
        Ok(WavWriter::new(file, spec))
    }

    /// Writes the RIFF WAVE header
    fn write_header(&mut self) -> io::Result<()> {
        // Useful links:
        // https://msdn.microsoft.com/en-us/library/ms713462.aspx
        // https://msdn.microsoft.com/en-us/library/ms713497.aspx

        let mut header = [0u8; 68];
        let spec = &self.spec;

        // Write the header in-memory first.
        {
            let mut buffer: io::Cursor<&mut [u8]> = io::Cursor::new(&mut header);
            try!(buffer.write_all("RIFF".as_bytes()));

            // Skip 4 bytes that will be filled with the file size afterwards.
            try!(buffer.write_le_u32(0));

            try!(buffer.write_all("WAVE".as_bytes()));
            try!(buffer.write_all("fmt ".as_bytes()));
            try!(buffer.write_le_u32(40)); // Size of the WAVE header chunk.

            // The following is based on the WAVEFORMATEXTENSIBLE struct as
            // documented on MSDN.

            // The field wFormatTag, value 1 means WAVE_FORMAT_PCM, but we use
            // the slightly more sophisticated WAVE_FORMAT_EXTENSIBLE.
            try!(buffer.write_le_u16(0xfffe));
            // The field nChannels.
            try!(buffer.write_le_u16(spec.channels));
            // The field nSamplesPerSec.
            try!(buffer.write_le_u32(spec.sample_rate));
            let bytes_per_sec = spec.sample_rate
                              * self.bytes_per_sample
                              * spec.channels as u32;
            // The field nAvgBytesPerSec;
            try!(buffer.write_le_u32(bytes_per_sec));
            // The field nBlockAlign. Block align * sample rate = bytes per sec.
            try!(buffer.write_le_u16((bytes_per_sec / spec.sample_rate) as u16));
            // The field wBitsPerSample. This is actually the size of the
            // container, so this is a multiple of 8.
            try!(buffer.write_le_u16(self.bytes_per_sample as u16 * 8));
            // The field cbSize, the number of remaining bytes in the struct.
            try!(buffer.write_le_u16(22));
            // The field wValidBitsPerSample, the real number of bits per sample.
            try!(buffer.write_le_u16(self.spec.bits_per_sample as u16));
            // The field dwChannelMask.
            // TODO: add the option to specify the channel mask. For now, use
            // the default assignment.
            try!(buffer.write_le_u32(channel_mask(self.spec.channels)));
            // The field SubFormat. We use KSDATAFORMAT_SUBTYPE_PCM. The
            // following GUIDS are defined:
            // - PCM:        00000001-0000-0010-8000-00aa00389b71
            // - IEEE_FLOAT: 00000003-0000-0010-8000-00aa00389b71
            // The byte order of a GUID is native for the first three sections,
            // which is assumed to be little endian, and big endian for the
            // last 8-byte section (which does contain a hyphen, for reasons
            // unknown to me).
            try!(buffer.write_all(&[0x01, 0x00, 0x00, 0x00,
                                    0x00, 0x00, 0x10, 0x00,
                                    0x80, 0x00, 0x00, 0xaa,
                                    0x00, 0x38, 0x9b, 0x71]));

            // So far the "fmt " section, now comes the "data" section. We will
            // only write the header here, actual data are the samples. The
            // number of bytes that this will take is not known at this point.
            // The 0 will be overwritten later.
            try!(buffer.write_all("data".as_bytes()));
            try!(buffer.write_le_u32(0));
        }

        // Then write the entire header at once.
        try!(self.writer.write_all(&header));

        Ok(())
    }

    /// Writes a single sample for one channel.
    ///
    /// WAVE interleaves channel data, so the channel that this writes the
    /// sample to depends on previous writes.
    pub fn write_sample<S: Sample>(&mut self, sample: S) -> io::Result<()> {
        if !self.wrote_header {
            try!(self.write_header());
            self.wrote_header = true;
        }

        // TODO: do we need bits per sample? Is the padding at the obvious side?
        try!(sample.write(&mut self.writer, self.bytes_per_sample));
        self.data_bytes_written += self.bytes_per_sample;
        Ok(())
    }

    /// Performs finalization. After calling this, the writer should be destructed.
    fn finalize_internal(&mut self) -> io::Result<()> {
        self.finalized = true;

        // Flush remaining samples via the BufWriter.
        try!(self.writer.flush());

        // Extract the underlying writer and rewind it to the start, to update
        // the header fields of which we now know the value.
        let mut writer = self.writer.get_mut();

        // The header minus magic and 32-bit filesize is 60 bytes long.
        let file_size = self.data_bytes_written + 60;
        try!(writer.seek(io::SeekFrom::Start(4)));
        try!(writer.write_le_u32(file_size));
        try!(writer.seek(io::SeekFrom::Start(64)));
        try!(writer.write_le_u32(self.data_bytes_written));

        Ok(())
    }

    /// Writes the parts of the WAVE format that require knowing all samples.
    ///
    /// This method must be called after all samples have been written. If it
    /// is not called, the destructor will finalize the file, but any IO errors
    /// that occur in the process cannot be observed in that manner.
    pub fn finalize(mut self) -> io::Result<()> {
        self.finalize_internal()
    }
}

impl<W> Drop for WavWriter<W> where W: io::Write + io::Seek {
    fn drop(&mut self) {
        // `finalize_internal` must be called only once. If that is done via
        // `finalize`, then this method is a no-op. If the user did not
        // finalize explicitly, then we should do it now. This can fail, but
        // drop should not panic, so a failure is ignored silently here.
        if !self.finalized {
            let _r = self.finalize_internal();
        }
    }
}

// TODO: Add benchmark for write speed.
<commit_msg>add skeleton for reading wav files<commit_after>// Hound -- A WAV encoding library in Rust
// Copyright (C) 2015 Ruud van Asseldonk
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3 of the License only.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

//! Hound, a WAV encoding library.
//!
//! TODO: Add some introductory text here.
//!
//! The following example renders a 440 Hz sine wave, and stores it as as a
//! mono wav file with a sample rate of 44.1 kHz and 16 bits per sample.
//!
//! ```
//! use std::fs;
//! use std::f32::consts::PI;
//! use std::i16;
//! use hound;
//!
//! let spec = hound::WavSpec {
//!     channels: 1,
//!     sample_rate: 44100,
//!     bits_per_sample: 16
//! };
//! // TODO: ensure that the type can be inferred.
//! let writer_res = hound::WavWriter::<fs::File>::create("sine.wav", spec);
//! let mut writer = writer_res.ok().unwrap();
//! for t in (0 .. 44100).map(|x| x as f32 / 44100.0) {
//!     let sample = (t * 440.0 * 2.0 * PI).sin();
//!     let amplitude: i16 = i16::MAX;
//!     writer.write_sample((sample * amplitude as f32) as i16).ok().unwrap();
//! }
//! writer.finalize().ok().unwrap();
//! ```

#![warn(missing_docs)]

use std::fs;
use std::io;
use std::io::Write;
use std::marker;
use std::path;

/// Extends the functionality of `io::Write` with additional methods.
///
/// The methods may be used on any type that implements `io::Write`.
trait WriteExt: io::Write {
    /// Writes an unsigned 16-bit integer in little endian format.
    fn write_le_u16(&mut self, x: u16) -> io::Result<()>;

    /// Writes an unsigned 32-bit integer in little endian format.
    fn write_le_u32(&mut self, x: u32) -> io::Result<()>;
}

impl<W> WriteExt for W where W: io::Write {

    fn write_le_u16(&mut self, x: u16) -> io::Result<()> {
        let mut buf = [0u8; 2];
        buf[0] = (x & 0xff) as u8;
        buf[1] = (x >> 8) as u8;
        self.write_all(&buf)
    }

    fn write_le_u32(&mut self, x: u32) -> io::Result<()> {
        let mut buf = [0u8; 4];
        buf[0] = ((x >> 00) & 0xff) as u8;
        buf[1] = ((x >> 08) & 0xff) as u8;
        buf[2] = ((x >> 16) & 0xff) as u8;
        buf[3] = ((x >> 24) & 0xff) as u8;
        self.write_all(&buf)
    }
}

/// A type that can be used to represent audio samples.
pub trait Sample {
    /// Writes the audio sample to the WAVE data section.
    fn write<W: io::Write>(self, writer: &mut W, bits: u32) -> io::Result<()>;
}

impl Sample for u16 {
    fn write<W: io::Write>(self, writer: &mut W, bits: u32) -> io::Result<()> {
        writer.write_le_u16(self)
        // TODO: take bits into account
    }
}

impl Sample for i16 {
    fn write<W: io::Write>(self, writer: &mut W, bits: u32) -> io::Result<()> {
        writer.write_le_u16(self as u16)
        // TODO: take bits into account
    }
}

/// Generates a bitmask with `channels` ones in the least significant bits.
fn channel_mask(channels: u16) -> u32 {
    (0 .. channels).map(|c| 1 << c).fold(0, |a, c| a | c)
}

#[test]
fn verify_channel_mask()
{
    assert_eq!(channel_mask(0), 0);
    assert_eq!(channel_mask(1), 1);
    assert_eq!(channel_mask(2), 3);
    assert_eq!(channel_mask(3), 7);
    assert_eq!(channel_mask(4), 15);
}

/// Specifies properties of the audio data.
#[derive(Clone, Copy)]
pub struct WavSpec {
    /// The number of channels.
    pub channels: u16,

    /// The number of samples per second.
    ///
    /// A common value is 44100, this is 44.1 kHz which is used for CD audio.
    pub sample_rate: u32,

    /// The number of bits per sample.
    ///
    /// A common value is 16 bits per sample, which is used for CD audio.
    pub bits_per_sample: u32
}

/// A writer that accepts samples and writes the WAVE format.
///
/// TODO: add example.
///
/// After all samples have been written, the file must be finalized. This can
/// be done by calling `finalize`. If `finalize` is not called, the file will
/// be finalized upon drop. However, finalization involves IO that may fail,
/// and without calling `finalize`, such a failure cannot be observed.
pub struct WavWriter<W> where W: io::Write + io::Seek {
    /// Specifies properties of the audio data.
    spec: WavSpec,

    /// The (container) bytes per sample. This is the bit rate / 8 rounded up.
    bytes_per_sample: u32,

    /// Whether the header has been written already.
    wrote_header: bool,

    /// The writer that will be written to.
    writer: io::BufWriter<W>,

    /// The number of bytes written to the data section.
    ///
    /// This is an `u32` because WAVE cannot accomodate more data.
    data_bytes_written: u32,

    /// Whether `finalize_internal` has been called.
    finalized: bool
}

impl<W> WavWriter<W> where W: io::Write + io::Seek {
    /// Creates a writer that writes the WAVE format to the underlying writer.
    ///
    /// The underlying writer is assumed to be at offset 0. `WavWriter` employs
    /// buffering internally to avoid too many `write` calls to the underlying
    /// writer.
    pub fn new(writer: W, spec: WavSpec) -> WavWriter<W> {
        WavWriter {
            spec: spec,
            bytes_per_sample: (spec.bits_per_sample as f32 / 8.0).ceil() as u32,
            wrote_header: false,
            writer: io::BufWriter::new(writer),
            data_bytes_written: 0,
            finalized: false
        }
    }

    /// Creates a writer that writes the WAVE format to a file.
    ///
    /// The file will be overwritten if it exists.
    pub fn create<P: AsRef<path::Path>>(filename: P, spec: WavSpec)
           -> io::Result<WavWriter<fs::File>> {
        let file = try!(fs::File::create(filename));
        Ok(WavWriter::new(file, spec))
    }

    /// Writes the RIFF WAVE header
    fn write_header(&mut self) -> io::Result<()> {
        // Useful links:
        // https://msdn.microsoft.com/en-us/library/ms713462.aspx
        // https://msdn.microsoft.com/en-us/library/ms713497.aspx

        let mut header = [0u8; 68];
        let spec = &self.spec;

        // Write the header in-memory first.
        {
            let mut buffer: io::Cursor<&mut [u8]> = io::Cursor::new(&mut header);
            try!(buffer.write_all("RIFF".as_bytes()));

            // Skip 4 bytes that will be filled with the file size afterwards.
            try!(buffer.write_le_u32(0));

            try!(buffer.write_all("WAVE".as_bytes()));
            try!(buffer.write_all("fmt ".as_bytes()));
            try!(buffer.write_le_u32(40)); // Size of the WAVE header chunk.

            // The following is based on the WAVEFORMATEXTENSIBLE struct as
            // documented on MSDN.

            // The field wFormatTag, value 1 means WAVE_FORMAT_PCM, but we use
            // the slightly more sophisticated WAVE_FORMAT_EXTENSIBLE.
            try!(buffer.write_le_u16(0xfffe));
            // The field nChannels.
            try!(buffer.write_le_u16(spec.channels));
            // The field nSamplesPerSec.
            try!(buffer.write_le_u32(spec.sample_rate));
            let bytes_per_sec = spec.sample_rate
                              * self.bytes_per_sample
                              * spec.channels as u32;
            // The field nAvgBytesPerSec;
            try!(buffer.write_le_u32(bytes_per_sec));
            // The field nBlockAlign. Block align * sample rate = bytes per sec.
            try!(buffer.write_le_u16((bytes_per_sec / spec.sample_rate) as u16));
            // The field wBitsPerSample. This is actually the size of the
            // container, so this is a multiple of 8.
            try!(buffer.write_le_u16(self.bytes_per_sample as u16 * 8));
            // The field cbSize, the number of remaining bytes in the struct.
            try!(buffer.write_le_u16(22));
            // The field wValidBitsPerSample, the real number of bits per sample.
            try!(buffer.write_le_u16(self.spec.bits_per_sample as u16));
            // The field dwChannelMask.
            // TODO: add the option to specify the channel mask. For now, use
            // the default assignment.
            try!(buffer.write_le_u32(channel_mask(self.spec.channels)));
            // The field SubFormat. We use KSDATAFORMAT_SUBTYPE_PCM. The
            // following GUIDS are defined:
            // - PCM:        00000001-0000-0010-8000-00aa00389b71
            // - IEEE_FLOAT: 00000003-0000-0010-8000-00aa00389b71
            // The byte order of a GUID is native for the first three sections,
            // which is assumed to be little endian, and big endian for the
            // last 8-byte section (which does contain a hyphen, for reasons
            // unknown to me).
            try!(buffer.write_all(&[0x01, 0x00, 0x00, 0x00,
                                    0x00, 0x00, 0x10, 0x00,
                                    0x80, 0x00, 0x00, 0xaa,
                                    0x00, 0x38, 0x9b, 0x71]));

            // So far the "fmt " section, now comes the "data" section. We will
            // only write the header here, actual data are the samples. The
            // number of bytes that this will take is not known at this point.
            // The 0 will be overwritten later.
            try!(buffer.write_all("data".as_bytes()));
            try!(buffer.write_le_u32(0));
        }

        // Then write the entire header at once.
        try!(self.writer.write_all(&header));

        Ok(())
    }

    /// Writes a single sample for one channel.
    ///
    /// WAVE interleaves channel data, so the channel that this writes the
    /// sample to depends on previous writes.
    pub fn write_sample<S: Sample>(&mut self, sample: S) -> io::Result<()> {
        if !self.wrote_header {
            try!(self.write_header());
            self.wrote_header = true;
        }

        // TODO: do we need bits per sample? Is the padding at the obvious side?
        try!(sample.write(&mut self.writer, self.bytes_per_sample));
        self.data_bytes_written += self.bytes_per_sample;
        Ok(())
    }

    /// Performs finalization. After calling this, the writer should be destructed.
    fn finalize_internal(&mut self) -> io::Result<()> {
        self.finalized = true;

        // Flush remaining samples via the BufWriter.
        try!(self.writer.flush());

        // Extract the underlying writer and rewind it to the start, to update
        // the header fields of which we now know the value.
        let mut writer = self.writer.get_mut();

        // The header minus magic and 32-bit filesize is 60 bytes long.
        let file_size = self.data_bytes_written + 60;
        try!(writer.seek(io::SeekFrom::Start(4)));
        try!(writer.write_le_u32(file_size));
        try!(writer.seek(io::SeekFrom::Start(64)));
        try!(writer.write_le_u32(self.data_bytes_written));

        Ok(())
    }

    /// Writes the parts of the WAVE format that require knowing all samples.
    ///
    /// This method must be called after all samples have been written. If it
    /// is not called, the destructor will finalize the file, but any IO errors
    /// that occur in the process cannot be observed in that manner.
    pub fn finalize(mut self) -> io::Result<()> {
        self.finalize_internal()
    }
}

impl<W> Drop for WavWriter<W> where W: io::Write + io::Seek {
    fn drop(&mut self) {
        // `finalize_internal` must be called only once. If that is done via
        // `finalize`, then this method is a no-op. If the user did not
        // finalize explicitly, then we should do it now. This can fail, but
        // drop should not panic, so a failure is ignored silently here.
        if !self.finalized {
            let _r = self.finalize_internal();
        }
    }
}

// TODO: Add benchmark for write speed.

/// A reader that reads the WAVE format from the underlying reader.
pub struct WavReader<R> {
    spec: WavSpec,
    reader: R
}

pub struct WavSamples<'wr, R, S> where R: 'wr {
    reader: &'wr mut WavReader<R>,
    phantom_sample: marker::PhantomData<S>
}

impl<R> WavReader<R> where R: io::Read {
    // TODO: define a custom error type to report ill-formed files.
    /// Attempts to create a reader that reads the WAVE format.
    ///
    /// The header is read immediately. Reading the data will be done on
    /// demand.
    pub fn new(reader: R) -> io::Result<WavReader<R>> {
        // TODO: try and read header, find spec.
        let spec = WavSpec {
            channels: 1,
            sample_rate: 44100,
            bits_per_sample: 16
        };
        let wav_reader = WavReader {
            spec: spec,
            reader: reader
        };

        Ok(wav_reader)
    }

    // TODO: Should this return by value instead? A reference is more consistent
    // with Claxon, but the type is only 80 bytes, barely larger than a pointer.
    // Is it worth the extra indirection? On the other hand, the indirection
    // is probably optimised away.
    /// Returns information about the WAVE file.
    pub fn spec(&self) -> &WavSpec {
        &self.spec
    }

    pub fn samples<'wr, S: Sample>(&'wr mut self) -> WavSamples<'wr, R, S> {
        WavSamples {
            reader: self,
            phantom_sample: marker::PhantomData
        }
    }
}
<|endoftext|>


<commit_before>// #![deny(warnings)]

extern crate fern;
#[macro_use]
extern crate log;
extern crate libc;
extern crate time;
extern crate xcb;
extern crate xcb_util;
extern crate xdg;

use std::cmp;
use std::rc::Rc;

pub mod cmd;
mod debug;
pub mod groups;
pub mod keys;
pub mod layout;
mod stack;
pub mod window;
pub mod x;

use groups::{Group, GroupBuilder};
use keys::{KeyCombo, KeyHandlers};
use layout::Layout;
use stack::Stack;
use x::{Connection, Event, StrutPartial, WindowId, WindowType};

/// Initializes a logger using the default configuration.
///
/// Outputs to stdout and $XDG_DATA/lanta/lanta.log by default.
/// You should feel free to initialize your own logger, instead of using this.
pub fn intiailize_logger() {
    let xdg_dirs = xdg::BaseDirectories::with_prefix("lanta")
        .expect("Could not create xdg BaseDirectories");
    let log_path = xdg_dirs
        .place_data_file("lanta.log")
        .expect("Could not create log file");

    let logger_config = fern::DispatchConfig {
        format: Box::new(
            |msg: &str, level: &log::LogLevel, _location: &log::LogLocation| {
                format!("[{}] [{}] {}", time::now().rfc3339(), level, msg)
            },
        ),
        output: vec![
            fern::OutputConfig::stdout(),
            fern::OutputConfig::file(&log_path),
        ],
        level: log::LogLevelFilter::Trace,
    };
    if let Err(e) = fern::init_global_logger(logger_config, log::LogLevelFilter::Trace) {
        panic!("Failed to initialize global logger: {}", e);
    }
}

#[macro_export]
macro_rules! keys {
    [ $( ([$( $mod:ident ),+], $key:ident, $cmd:expr) ),+ $(,)*] => (
        vec![
            $( (vec![$( $mod ),+], ::x11::keysym::$key, $cmd) ),+
        ]
    )
}

#[macro_export]
macro_rules! groups {
    { $keys:ident, [$(( [$( $modkey:ident ),+], $key:ident, $name:expr, $layout:expr )),+ $(,)*] }  => {{
        $keys.extend(keys![
            $(
                ([$($modkey),+], $key, cmd::lazy::switch_group($name))
            ),+
        ]);
        vec![
            $(
                GroupBuilder::new($name, $layout)
            ),+
        ]
    }}
}

#[derive(Clone, Copy, Debug, Default)]
pub struct Viewport {
    pub x: u32,
    pub y: u32,
    pub width: u32,
    pub height: u32,
}

struct Dock {
    window_id: WindowId,
    strut_partial: Option<StrutPartial>,
}

#[derive(Default)]
struct Docks {
    vec: Vec<Dock>,
}

impl Docks {
    pub fn add(&mut self, conn: &Connection, window_id: WindowId) {
        let strut_partial = conn.get_strut_partial(&window_id);
        self.vec.push(Dock {
            window_id,
            strut_partial,
        });
    }

    pub fn remove(&mut self, window_id: &WindowId) {
        self.vec.retain(|d| &d.window_id != window_id);
    }

    /// Figure out the usable area of the screen based on the STRUT_PARTIAL of
    /// all docks.
    pub fn viewport(&self, screen_width: u32, screen_height: u32) -> Viewport {
        let (left, right, top, bottom) = self.vec
            .iter()
            .filter_map(|o| o.strut_partial.as_ref())
            .fold((0, 0, 0, 0), |(left, right, top, bottom), s| {
                // We don't bother looking at the start/end members of the
                // StrutPartial - treating it more like a Strut.
                (
                    cmp::max(left, s.left()),
                    cmp::max(right, s.right()),
                    cmp::max(top, s.top()),
                    cmp::max(bottom, s.bottom()),
                )
            });
        let viewport = Viewport {
            x: left,
            y: top,
            width: screen_width - left - right,
            height: screen_height - top - bottom,
        };
        debug!("Calculated Viewport as {:?}", viewport);
        viewport
    }
}

pub struct Lanta {
    connection: Rc<Connection>,
    keys: KeyHandlers,
    groups: Stack<Group>,
    docks: Docks,
}

impl Lanta {
    pub fn new<K>(
        keys: K,
        groups: Vec<GroupBuilder>,
        layouts: Vec<Box<Layout>>,
    ) -> Result<Self, String>
    where
        K: Into<KeyHandlers>,
    {
        let keys = keys.into();
        let connection = Rc::new(Connection::connect()?);
        connection.install_as_wm(&keys)?;

        let groups = Stack::from(
            groups
                .into_iter()
                .map(|group: GroupBuilder| {
                    group.build(connection.clone(), layouts.clone())
                })
                .collect::<Vec<Group>>(),
        );

        let mut wm = Lanta {
            connection: connection.clone(),
            keys: keys,
            groups: groups,
            docks: Docks::default(),
        };

        // Learn about existing top-level windows.
        let existing_windows = connection.top_level_windows().unwrap();
        for window in existing_windows {
            wm.manage_window(window);
        }
        let viewport = wm.viewport();
        wm.group_mut().activate(viewport);
        wm.connection.update_ewmh_desktops(&wm.groups);

        Ok(wm)
    }

    fn viewport(&self) -> Viewport {
        let (width, height) = self.connection
            .get_window_geometry(self.connection.root_window_id());
        self.docks.viewport(width, height)
    }

    pub fn group(&self) -> &Group {
        self.groups.focused().expect("No active group!")
    }

    pub fn group_mut(&mut self) -> &mut Group {
        self.groups.focused_mut().expect("No active group!")
    }

    pub fn switch_group<'a, S>(&'a mut self, name: S)
    where
        S: Into<&'a str>,
    {
        let name = name.into();

        // If we're already on this group, do nothing.
        if self.group().name() == name {
            return;
        }

        self.group_mut().deactivate();
        self.groups.focus(|group| group.name() == name);
        let viewport = self.viewport();
        self.group_mut().activate(viewport);
        self.connection.update_ewmh_desktops(&self.groups);
    }

    /// Move the focused window from the active group to another named group.
    ///
    /// If the other named group does not exist, then the window is
    /// (unfortunately) lost.
    pub fn move_focused_to_group<'a, S>(&'a mut self, name: S)
    where
        S: Into<&'a str>,
    {
        let name = name.into();

        // If the group is currently active, then do nothing. This avoids flicker as we
        // unmap/remap.
        if name == self.group().name() {
            return;
        }

        let removed = self.group_mut().remove_focused();
        let new_group_opt = self.groups.iter_mut().find(|group| group.name() == name);
        match new_group_opt {
            Some(new_group) => {
                removed.map(|window| new_group.add_window(window));
            }
            None => {
                // It would be nice to put the window back in its group (or avoid taking it out
                // of its group until we've checked the new group exists), but it's difficult
                // to do this while keeping the borrow checker happy.
                error!("Moved window to non-existent group: {}", name);
            }
        }
    }

    pub fn manage_window(&mut self, window_id: WindowId) {
        debug!("Managing window: {}", window_id);

        // If we are already managing the window, then do nothing.
        // This may result in some undesireableness, in that we don't move
        // it to the currently focused group when it's remapped, but it shouldn't
        // be too bad.
        let already_managed = self.groups.iter().any(|g| g.contains(&window_id));
        if already_managed {
            warn!(
                "Asked to manage window that's already managed: {}",
                window_id
            );
            return;
        }

        let window_types = self.connection.get_window_types(&window_id);
        let dock = window_types.contains(&WindowType::Dock);

        self.connection
            .enable_window_key_events(&window_id, &self.keys);

        if dock {
            self.docks.add(&self.connection, window_id);
            let viewport = self.viewport();
            self.group_mut().update_viewport(viewport);
        } else {
            self.connection.enable_window_tracking(&window_id);
            self.group_mut().add_window(window_id);
        }
    }

    pub fn unmanage_window(&mut self, window_id: WindowId) {
        debug!("Unmanaging window: {}", window_id);

        // Remove the window from whichever Group it is in. Special case for
        // docks which aren't in any group.
        self.groups
            .iter_mut()
            .find(|group| group.contains(&window_id))
            .map(|group| group.remove_window(&window_id));
        self.docks.remove(&window_id);

        // The viewport may have changed.
        let viewport = self.viewport();
        self.group_mut().update_viewport(viewport);
    }

    pub fn run(mut self) {
        info!("Started WM, entering event loop.");
        let event_loop_connection = self.connection.clone();
        let event_loop = event_loop_connection.get_event_loop();
        for event in event_loop {
            match event {
                Event::MapRequest(window_id) => self.on_map_request(window_id),
                Event::UnmapNotify(window_id) => self.on_unmap_notify(window_id),
                Event::DestroyNotify(window_id) => self.on_destroy_notify(window_id),
                Event::KeyPress(key) => self.on_key_press(key),
                Event::EnterNotify(window_id) => self.on_enter_notify(window_id),
            }
        }
        info!("Event loop exiting");
    }

    fn on_map_request(&mut self, window_id: WindowId) {
        self.connection.map_window(&window_id);
        self.manage_window(window_id);
    }

    fn on_unmap_notify(&mut self, window_id: WindowId) {
        // We only receive an unmap notify event when the window is actually
        // unmapped by its application. When our layouts unmap windows, they
        // (should) do it by disabling event tracking first.
        self.unmanage_window(window_id);
    }

    fn on_destroy_notify(&mut self, window_id: WindowId) {
        self.unmanage_window(window_id);
    }

    fn on_key_press(&mut self, key: KeyCombo) {
        self.keys.get(&key).map(move |handler| (handler)(self));
    }

    fn on_enter_notify(&mut self, window_id: WindowId) {
        self.group_mut().focus(&window_id);
    }
}
<commit_msg>Break up overly long line.<commit_after>// #![deny(warnings)]

extern crate fern;
#[macro_use]
extern crate log;
extern crate libc;
extern crate time;
extern crate xcb;
extern crate xcb_util;
extern crate xdg;

use std::cmp;
use std::rc::Rc;

pub mod cmd;
mod debug;
pub mod groups;
pub mod keys;
pub mod layout;
mod stack;
pub mod window;
pub mod x;

use groups::{Group, GroupBuilder};
use keys::{KeyCombo, KeyHandlers};
use layout::Layout;
use stack::Stack;
use x::{Connection, Event, StrutPartial, WindowId, WindowType};

/// Initializes a logger using the default configuration.
///
/// Outputs to stdout and $XDG_DATA/lanta/lanta.log by default.
/// You should feel free to initialize your own logger, instead of using this.
pub fn intiailize_logger() {
    let xdg_dirs = xdg::BaseDirectories::with_prefix("lanta")
        .expect("Could not create xdg BaseDirectories");
    let log_path = xdg_dirs
        .place_data_file("lanta.log")
        .expect("Could not create log file");

    let logger_config = fern::DispatchConfig {
        format: Box::new(
            |msg: &str, level: &log::LogLevel, _location: &log::LogLocation| {
                format!("[{}] [{}] {}", time::now().rfc3339(), level, msg)
            },
        ),
        output: vec![
            fern::OutputConfig::stdout(),
            fern::OutputConfig::file(&log_path),
        ],
        level: log::LogLevelFilter::Trace,
    };
    if let Err(e) = fern::init_global_logger(logger_config, log::LogLevelFilter::Trace) {
        panic!("Failed to initialize global logger: {}", e);
    }
}

#[macro_export]
macro_rules! keys {
    [ $( ([$( $mod:ident ),+], $key:ident, $cmd:expr) ),+ $(,)*] => (
        vec![
            $( (vec![$( $mod ),+], ::x11::keysym::$key, $cmd) ),+
        ]
    )
}

#[macro_export]
macro_rules! groups {
    {
        $keys:ident,
        [
            $(( [$( $modkey:ident ),+], $key:ident, $name:expr, $layout:expr )),+
            $(,)*
        ]
    }  => {{
        $keys.extend(keys![
            $(
                ([$($modkey),+], $key, cmd::lazy::switch_group($name))
            ),+
        ]);
        vec![
            $(
                GroupBuilder::new($name, $layout)
            ),+
        ]
    }}
}

#[derive(Clone, Copy, Debug, Default)]
pub struct Viewport {
    pub x: u32,
    pub y: u32,
    pub width: u32,
    pub height: u32,
}

struct Dock {
    window_id: WindowId,
    strut_partial: Option<StrutPartial>,
}

#[derive(Default)]
struct Docks {
    vec: Vec<Dock>,
}

impl Docks {
    pub fn add(&mut self, conn: &Connection, window_id: WindowId) {
        let strut_partial = conn.get_strut_partial(&window_id);
        self.vec.push(Dock {
            window_id,
            strut_partial,
        });
    }

    pub fn remove(&mut self, window_id: &WindowId) {
        self.vec.retain(|d| &d.window_id != window_id);
    }

    /// Figure out the usable area of the screen based on the STRUT_PARTIAL of
    /// all docks.
    pub fn viewport(&self, screen_width: u32, screen_height: u32) -> Viewport {
        let (left, right, top, bottom) = self.vec
            .iter()
            .filter_map(|o| o.strut_partial.as_ref())
            .fold((0, 0, 0, 0), |(left, right, top, bottom), s| {
                // We don't bother looking at the start/end members of the
                // StrutPartial - treating it more like a Strut.
                (
                    cmp::max(left, s.left()),
                    cmp::max(right, s.right()),
                    cmp::max(top, s.top()),
                    cmp::max(bottom, s.bottom()),
                )
            });
        let viewport = Viewport {
            x: left,
            y: top,
            width: screen_width - left - right,
            height: screen_height - top - bottom,
        };
        debug!("Calculated Viewport as {:?}", viewport);
        viewport
    }
}

pub struct Lanta {
    connection: Rc<Connection>,
    keys: KeyHandlers,
    groups: Stack<Group>,
    docks: Docks,
}

impl Lanta {
    pub fn new<K>(
        keys: K,
        groups: Vec<GroupBuilder>,
        layouts: Vec<Box<Layout>>,
    ) -> Result<Self, String>
    where
        K: Into<KeyHandlers>,
    {
        let keys = keys.into();
        let connection = Rc::new(Connection::connect()?);
        connection.install_as_wm(&keys)?;

        let groups = Stack::from(
            groups
                .into_iter()
                .map(|group: GroupBuilder| {
                    group.build(connection.clone(), layouts.clone())
                })
                .collect::<Vec<Group>>(),
        );

        let mut wm = Lanta {
            connection: connection.clone(),
            keys: keys,
            groups: groups,
            docks: Docks::default(),
        };

        // Learn about existing top-level windows.
        let existing_windows = connection.top_level_windows().unwrap();
        for window in existing_windows {
            wm.manage_window(window);
        }
        let viewport = wm.viewport();
        wm.group_mut().activate(viewport);
        wm.connection.update_ewmh_desktops(&wm.groups);

        Ok(wm)
    }

    fn viewport(&self) -> Viewport {
        let (width, height) = self.connection
            .get_window_geometry(self.connection.root_window_id());
        self.docks.viewport(width, height)
    }

    pub fn group(&self) -> &Group {
        self.groups.focused().expect("No active group!")
    }

    pub fn group_mut(&mut self) -> &mut Group {
        self.groups.focused_mut().expect("No active group!")
    }

    pub fn switch_group<'a, S>(&'a mut self, name: S)
    where
        S: Into<&'a str>,
    {
        let name = name.into();

        // If we're already on this group, do nothing.
        if self.group().name() == name {
            return;
        }

        self.group_mut().deactivate();
        self.groups.focus(|group| group.name() == name);
        let viewport = self.viewport();
        self.group_mut().activate(viewport);
        self.connection.update_ewmh_desktops(&self.groups);
    }

    /// Move the focused window from the active group to another named group.
    ///
    /// If the other named group does not exist, then the window is
    /// (unfortunately) lost.
    pub fn move_focused_to_group<'a, S>(&'a mut self, name: S)
    where
        S: Into<&'a str>,
    {
        let name = name.into();

        // If the group is currently active, then do nothing. This avoids flicker as we
        // unmap/remap.
        if name == self.group().name() {
            return;
        }

        let removed = self.group_mut().remove_focused();
        let new_group_opt = self.groups.iter_mut().find(|group| group.name() == name);
        match new_group_opt {
            Some(new_group) => {
                removed.map(|window| new_group.add_window(window));
            }
            None => {
                // It would be nice to put the window back in its group (or avoid taking it out
                // of its group until we've checked the new group exists), but it's difficult
                // to do this while keeping the borrow checker happy.
                error!("Moved window to non-existent group: {}", name);
            }
        }
    }

    pub fn manage_window(&mut self, window_id: WindowId) {
        debug!("Managing window: {}", window_id);

        // If we are already managing the window, then do nothing.
        // This may result in some undesireableness, in that we don't move
        // it to the currently focused group when it's remapped, but it shouldn't
        // be too bad.
        let already_managed = self.groups.iter().any(|g| g.contains(&window_id));
        if already_managed {
            warn!(
                "Asked to manage window that's already managed: {}",
                window_id
            );
            return;
        }

        let window_types = self.connection.get_window_types(&window_id);
        let dock = window_types.contains(&WindowType::Dock);

        self.connection
            .enable_window_key_events(&window_id, &self.keys);

        if dock {
            self.docks.add(&self.connection, window_id);
            let viewport = self.viewport();
            self.group_mut().update_viewport(viewport);
        } else {
            self.connection.enable_window_tracking(&window_id);
            self.group_mut().add_window(window_id);
        }
    }

    pub fn unmanage_window(&mut self, window_id: WindowId) {
        debug!("Unmanaging window: {}", window_id);

        // Remove the window from whichever Group it is in. Special case for
        // docks which aren't in any group.
        self.groups
            .iter_mut()
            .find(|group| group.contains(&window_id))
            .map(|group| group.remove_window(&window_id));
        self.docks.remove(&window_id);

        // The viewport may have changed.
        let viewport = self.viewport();
        self.group_mut().update_viewport(viewport);
    }

    pub fn run(mut self) {
        info!("Started WM, entering event loop.");
        let event_loop_connection = self.connection.clone();
        let event_loop = event_loop_connection.get_event_loop();
        for event in event_loop {
            match event {
                Event::MapRequest(window_id) => self.on_map_request(window_id),
                Event::UnmapNotify(window_id) => self.on_unmap_notify(window_id),
                Event::DestroyNotify(window_id) => self.on_destroy_notify(window_id),
                Event::KeyPress(key) => self.on_key_press(key),
                Event::EnterNotify(window_id) => self.on_enter_notify(window_id),
            }
        }
        info!("Event loop exiting");
    }

    fn on_map_request(&mut self, window_id: WindowId) {
        self.connection.map_window(&window_id);
        self.manage_window(window_id);
    }

    fn on_unmap_notify(&mut self, window_id: WindowId) {
        // We only receive an unmap notify event when the window is actually
        // unmapped by its application. When our layouts unmap windows, they
        // (should) do it by disabling event tracking first.
        self.unmanage_window(window_id);
    }

    fn on_destroy_notify(&mut self, window_id: WindowId) {
        self.unmanage_window(window_id);
    }

    fn on_key_press(&mut self, key: KeyCombo) {
        self.keys.get(&key).map(move |handler| (handler)(self));
    }

    fn on_enter_notify(&mut self, window_id: WindowId) {
        self.group_mut().focus(&window_id);
    }
}
<|endoftext|>


<commit_before>#![feature(convert,custom_derive,plugin)]
#![plugin(serde_macros)]
extern crate hyper;
extern crate serde;
extern crate serde_json;

use std::io::Read;

use hyper::Client;
use hyper::client::IntoUrl;
use hyper::client::Response;
use hyper::net::HttpsConnector;
use hyper::net::Openssl;
use serde_json::Value;

pub struct LxdServer {
    url: String,
    cert_location: String,
    key_location: String,
}

impl LxdServer {

    pub fn new(url: &str, cert_location: &str, key_location: &str) -> LxdServer {
        LxdServer {
            url: url.to_string(),
            cert_location: cert_location.to_string(),
            key_location: key_location.to_string()
        }
    }

    fn get_client(&self) -> Client {
        let connector = HttpsConnector::new(
            Openssl::with_cert_and_key(&self.cert_location, &self.key_location).unwrap());
        Client::with_connector(connector)
    }

    pub fn get(&self, relative_url: &str) -> Response {
        let client = self.get_client();
        let url = match (self.url.clone() + relative_url).into_url() {
            Err(why) => panic!("{:?}", why),
            Ok(url) => url,
        };
        let unsent_response = client.get(url);
        match unsent_response.send() {
            Err(why) => panic!("{:?}", why),
            Ok(response) => response,
        }
    }
}

#[derive(Deserialize)]
pub struct ContainerIP {
    pub address: String,
    pub host_veth: String,
    pub interface: String,
    pub protocol: String,
}

pub struct ContainerStatus {
    pub status: String,
    pub ips: Vec<ContainerIP>,
}

pub struct Container {
    pub name: String,
    pub status: ContainerStatus,
    pub ephemeral: bool,
    pub snapshot_urls: Vec<String>,
//    pub ipv4: String,
//    pub ipv6: String,
}

impl Container {

    pub fn from_json(json: &Value) -> Container {
        let value_for_path = |path: &[&str]| {
            match json.find_path(&path) {
                Some(value) => { value }
                None => panic!("Couldn't find {} in JSON", path.join("."))
            }
        };
        let get_array_from_json = |path: &[&str]| {
            match value_for_path(&path).as_array() {
                Some(array) => { array }
                None => panic!("Couldn't find array at {}", path.join("."))
            }
        };
        let get_boolean_from_json = |path: &[&str]| {
            match value_for_path(&path).as_boolean() {
                Some(boolean) => { boolean }
                None => panic!("Couldn't find boolean at {}", path.join("."))
            }
        };
        let get_string_from_json = |path: &[&str]| {
            match value_for_path(&path).as_string() {
                Some(string) => { string.to_string() }
                None => panic!("Couldn't find string at {}", path.join("."))
            }
        };
        Container {
            ephemeral: get_boolean_from_json(&["state", "ephemeral"]),
            name: get_string_from_json(&["state", "name"]),
            status: ContainerStatus {
                status: get_string_from_json(&["state", "status", "status"]),
                ips: get_array_from_json(&["state", "status", "ips"]).iter().map(
                    |ip_value| -> ContainerIP {
                        match serde_json::from_value(ip_value.clone()) {
                            Ok(container_ip) => container_ip,
                            Err(_) => panic!("Couldn't parse a ContainerIP for {:?}", ip_value)
                        }}).collect(),
            },
            snapshot_urls: match json.find_path(&["snaps"]).unwrap().as_array() {
                Some(array) => {
                    array.iter().map(|x| {x.as_string().unwrap().to_string()}).collect()
                }
                None => {
                    vec![]
                }
            }
        }
    }
}

fn response_to_value(response: &mut Response) -> Value {
    let mut body = String::new();
    response.read_to_string(&mut body).unwrap();
    serde_json::from_str(body.as_str()).unwrap()
}

pub fn list_containers() ->  Vec<Container> {
    let server = LxdServer::new(
        "https://104.155.75.254:8443",
        "/home/daniel/.config/lxc/client.crt",
        "/home/daniel/.config/lxc/client.key"
    );
    let mut response = server.get("/1.0/containers?recursion=1");
    let payload = response_to_value(&mut response);
    let container_values = payload.find("metadata").unwrap().as_array().unwrap();
    let mut containers = Vec::new();
    for container_value in container_values {
        let container = Container::from_json(container_value);
        containers.push(container);
    }
    containers
}
<commit_msg>Map FTW.<commit_after>#![feature(convert,custom_derive,plugin)]
#![plugin(serde_macros)]
extern crate hyper;
extern crate serde;
extern crate serde_json;

use std::io::Read;

use hyper::Client;
use hyper::client::IntoUrl;
use hyper::client::Response;
use hyper::net::HttpsConnector;
use hyper::net::Openssl;
use serde_json::Value;

pub struct LxdServer {
    url: String,
    cert_location: String,
    key_location: String,
}

impl LxdServer {

    pub fn new(url: &str, cert_location: &str, key_location: &str) -> LxdServer {
        LxdServer {
            url: url.to_string(),
            cert_location: cert_location.to_string(),
            key_location: key_location.to_string()
        }
    }

    fn get_client(&self) -> Client {
        let connector = HttpsConnector::new(
            Openssl::with_cert_and_key(&self.cert_location, &self.key_location).unwrap());
        Client::with_connector(connector)
    }

    pub fn get(&self, relative_url: &str) -> Response {
        let client = self.get_client();
        let url = match (self.url.clone() + relative_url).into_url() {
            Err(why) => panic!("{:?}", why),
            Ok(url) => url,
        };
        let unsent_response = client.get(url);
        match unsent_response.send() {
            Err(why) => panic!("{:?}", why),
            Ok(response) => response,
        }
    }
}

#[derive(Deserialize)]
pub struct ContainerIP {
    pub address: String,
    pub host_veth: String,
    pub interface: String,
    pub protocol: String,
}

pub struct ContainerStatus {
    pub status: String,
    pub ips: Vec<ContainerIP>,
}

pub struct Container {
    pub name: String,
    pub status: ContainerStatus,
    pub ephemeral: bool,
    pub snapshot_urls: Vec<String>,
//    pub ipv4: String,
//    pub ipv6: String,
}

impl Container {

    pub fn from_json(json: &Value) -> Container {
        let value_for_path = |path: &[&str]| {
            match json.find_path(&path) {
                Some(value) => { value }
                None => panic!("Couldn't find {} in JSON", path.join("."))
            }
        };
        let get_array_from_json = |path: &[&str]| {
            match value_for_path(&path).as_array() {
                Some(array) => { array }
                None => panic!("Couldn't find array at {}", path.join("."))
            }
        };
        let get_boolean_from_json = |path: &[&str]| {
            match value_for_path(&path).as_boolean() {
                Some(boolean) => { boolean }
                None => panic!("Couldn't find boolean at {}", path.join("."))
            }
        };
        let get_string_from_json = |path: &[&str]| {
            match value_for_path(&path).as_string() {
                Some(string) => { string.to_string() }
                None => panic!("Couldn't find string at {}", path.join("."))
            }
        };
        Container {
            ephemeral: get_boolean_from_json(&["state", "ephemeral"]),
            name: get_string_from_json(&["state", "name"]),
            status: ContainerStatus {
                status: get_string_from_json(&["state", "status", "status"]),
                ips: get_array_from_json(&["state", "status", "ips"]).iter().map(
                    |ip_value| -> ContainerIP {
                        match serde_json::from_value(ip_value.clone()) {
                            Ok(container_ip) => container_ip,
                            Err(_) => panic!("Couldn't parse a ContainerIP for {:?}", ip_value)
                        }}).collect(),
            },
            snapshot_urls: match json.find_path(&["snaps"]).unwrap().as_array() {
                Some(array) => {
                    array.iter().map(|x| {x.as_string().unwrap().to_string()}).collect()
                }
                None => {
                    vec![]
                }
            }
        }
    }
}

fn response_to_value(response: &mut Response) -> Value {
    let mut body = String::new();
    response.read_to_string(&mut body).unwrap();
    serde_json::from_str(body.as_str()).unwrap()
}

pub fn list_containers() ->  Vec<Container> {
    let server = LxdServer::new(
        "https://104.155.75.254:8443",
        "/home/daniel/.config/lxc/client.crt",
        "/home/daniel/.config/lxc/client.key"
    );
    let mut response = server.get("/1.0/containers?recursion=1");
    let payload = response_to_value(&mut response);
    let container_values = payload.find("metadata").unwrap().as_array().unwrap();
    container_values.iter().map(|v| { Container::from_json(v) }).collect()
}
<|endoftext|>


<commit_before>use std::error::Error;
use std::str::FromStr;

pub fn puzzle(input: &str) -> i32 {
    let instructions: Result<Vec<_>, _> = input.lines().map(|line| line.trim().parse()).collect();
    let instructions = instructions.expect("couldn't parse instructions");
    let mut machine = Machine::new();
    machine.execute(&instructions);
    machine.value_of(Register::A)
}

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Instruction {
    Copy(FromLocation, Register),
    Increment(Register),
    Decrement(Register),
    JumpNonZero(FromLocation, i32),
}

impl FromStr for Instruction {
    type Err = Box<Error>;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        println!("parsing instruction: {}", s);

        let mut parts = s.split_whitespace();
        match parts.next() {
            Some("cpy") => Ok(
                Instruction::Copy(
                    parts.next().ok_or("no from location")?.parse()?,
                    parts.next().ok_or("no register")?.parse()?
                )
            ),
            Some("inc") => Ok(
                Instruction::Increment(
                    parts.next().ok_or("no register")?.parse()?
                )
            ),
            Some("dec") => Ok(
                Instruction::Decrement(
                    parts.next().ok_or("no register")?.parse()?
                )
            ),
            Some("jnz") => Ok(
                Instruction::JumpNonZero(
                    parts.next().ok_or("no register")?.parse()?,
                    parts.next().ok_or("no offset")?.parse()?
                )
            ),
            Some(other) => Err(format!("Unknown instruction: {}", other).into()),
            None => Err("Instruction missing".into()),
        }
    }
}

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum FromLocation {
    Integer(i32),
    Register(Register),
}

impl FromStr for FromLocation {
    type Err = Box<Error>;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        s.parse().map(FromLocation::Integer)
                 .or_else(|_| s.parse().map(FromLocation::Register))
    }
}

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Register {
    A,
    B,
    C,
    D,
}

impl FromStr for Register {
    type Err = Box<Error>;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        use Register::*;
        match s {
            "a" => Ok(A),
            "b" => Ok(B),
            "c" => Ok(C),
            "d" => Ok(D),
            other => Err(format!("Unknown register {}", other).into()),
        }
    }
}

pub struct Machine {
    registers: [i32; 4],
}

impl Machine {
    pub fn new() -> Machine {
        Machine {
            registers: [0; 4],
        }
    }

    fn execute(&mut self, instructions: &[Instruction]) {
        let mut program_counter = 0;
        while let Some(&instruction) = instructions.get(program_counter) {
            // println!("registers = {:?}", self.registers);
            // println!("executing instruction at {}: {:?}", program_counter, instruction);

            let mut jump = 0;
            match instruction {
                Instruction::Copy(
                    FromLocation::Integer(i),
                    to
                ) => {
                    let index = Machine::register_index(to);
                    self.registers[index] = i;
                },
                Instruction::Copy(FromLocation::Register(from), to) => {
                    let index_from = Machine::register_index(from);
                    let index_to = Machine::register_index(to);
                    self.registers[index_to] = self.registers[index_from];
                },
                Instruction::Increment(register) => {
                    let index = Machine::register_index(register);
                    self.registers[index] += 1;
                },
                Instruction::Decrement(register) => {
                    let index = Machine::register_index(register);
                    self.registers[index] -= 1;
                },
                Instruction::JumpNonZero(FromLocation::Integer(i), offset) => {
                    if i != 0 {
                        jump = offset - 1;
                    }
                },
                Instruction::JumpNonZero(FromLocation::Register(register), offset) => {
                    let index = Machine::register_index(register);
                    if self.registers[index] != 0 {
                        jump = offset - 1;
                    }
                },
            }
            let next_pc = program_counter as i32 + 1 + jump;
            if next_pc < 0 {
                panic!("Jumped too far back!")
            }
            program_counter = next_pc as usize;
        }
    }

    fn register_index(register: Register) -> usize {
        use Register::*;
        match register {
            A => 0,
            B => 1,
            C => 2,
            D => 3,
        }
    }

    fn value_of(&self, register: Register) -> i32 {
        let index = Machine::register_index(register);
        self.registers[index]
    }
}

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

    #[test]
    fn parse_cpy() {
        let instr: Instruction = "cpy 41 a".parse().expect("couldn't parse");
        assert_eq!(
            instr,
            Instruction::Copy(FromLocation::Integer(41), Register::A)
        );
    }

    #[test]
    fn parse_cpy_register() {
        let instr: Instruction = "cpy a c".parse().expect("couldn't parse");
        assert_eq!(
            instr,
            Instruction::Copy(
                FromLocation::Register(Register::A),
                Register::C,
            )
        );
    }

    #[test]
    fn parse_inc() {
        let instr: Instruction = "inc a".parse().expect("couldn't parse");
        assert_eq!(
            instr,
            Instruction::Increment(Register::A)
        );
    }

    #[test]
    fn parse_dec() {
        let instr: Instruction = "dec a".parse().expect("couldn't parse");
        assert_eq!(
            instr,
            Instruction::Decrement(Register::A)
        );
    }

    #[test]
    fn parse_jnz() {
        let instr: Instruction = "jnz a 2".parse().expect("couldn't parse");
        assert_eq!(
            instr,
            Instruction::JumpNonZero(FromLocation::Register(Register::A), 2)
        );
    }

    #[test]
    fn execute_copy() {
        let mut machine = Machine::new();
        machine.execute(
            &[
                Instruction::Copy(
                    FromLocation::Integer(5),
                    Register::A,
                ),
            ]
        );

        assert_eq!(machine.value_of(Register::A), 5);
    }

    #[test]
    fn execute_copy_register() {
        let mut machine = Machine::new();
        machine.execute(
            &[
                Instruction::Copy(
                    FromLocation::Integer(5),
                    Register::A,
                ),
                Instruction::Copy(
                    FromLocation::Register(Register::A),
                    Register::B,
                ),
            ]
        );

        assert_eq!(machine.value_of(Register::B), 5);
    }

    #[test]
    fn execute_increment() {
        let mut machine = Machine::new();
        machine.execute(
            &[
                Instruction::Increment(
                    Register::C,
                ),
            ]
        );

        assert_eq!(machine.value_of(Register::C), 1);
    }

    #[test]
    fn execute_decrement() {
        let mut machine = Machine::new();
        machine.execute(
            &[
                Instruction::Decrement(
                    Register::D,
                ),
            ]
        );

        assert_eq!(machine.value_of(Register::D), -1);
    }

    #[test]
    fn execute_jnz_when_zero() {
        let mut machine = Machine::new();
        machine.execute(
            &[
                // D should be 0, jump doesn't happen
                Instruction::JumpNonZero(FromLocation::Register(Register::D), 2),
                Instruction::Increment(
                    Register::A,
                ),
            ]
        );

        assert_eq!(machine.value_of(Register::A), 1);
    }

    #[test]
    fn execute_jnz_when_non_zero() {
        let mut machine = Machine::new();
        machine.execute(
            &[
                Instruction::Decrement(
                    Register::D,
                ),
                Instruction::JumpNonZero(FromLocation::Register(Register::D), 2),
                Instruction::Increment(
                    Register::A,
                ),
            ]
        );

        assert_eq!(machine.value_of(Register::A), 0);
    }

    #[test]
    fn sample() {
        let input = "cpy 41 a
            inc a
            inc a
            dec a
            jnz a 2
            dec a";
        assert_eq!(puzzle(input), 42);
    }
}
<commit_msg>Change for part 2<commit_after>use std::error::Error;
use std::str::FromStr;

pub fn puzzle(input: &str) -> i32 {
    let instructions: Result<Vec<_>, _> = input.lines().map(|line| line.trim().parse()).collect();
    let instructions = instructions.expect("couldn't parse instructions");
    let mut machine = Machine::new();
    machine.execute(&instructions);
    machine.value_of(Register::A)
}

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Instruction {
    Copy(FromLocation, Register),
    Increment(Register),
    Decrement(Register),
    JumpNonZero(FromLocation, i32),
}

impl FromStr for Instruction {
    type Err = Box<Error>;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        println!("parsing instruction: {}", s);

        let mut parts = s.split_whitespace();
        match parts.next() {
            Some("cpy") => Ok(
                Instruction::Copy(
                    parts.next().ok_or("no from location")?.parse()?,
                    parts.next().ok_or("no register")?.parse()?
                )
            ),
            Some("inc") => Ok(
                Instruction::Increment(
                    parts.next().ok_or("no register")?.parse()?
                )
            ),
            Some("dec") => Ok(
                Instruction::Decrement(
                    parts.next().ok_or("no register")?.parse()?
                )
            ),
            Some("jnz") => Ok(
                Instruction::JumpNonZero(
                    parts.next().ok_or("no register")?.parse()?,
                    parts.next().ok_or("no offset")?.parse()?
                )
            ),
            Some(other) => Err(format!("Unknown instruction: {}", other).into()),
            None => Err("Instruction missing".into()),
        }
    }
}

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum FromLocation {
    Integer(i32),
    Register(Register),
}

impl FromStr for FromLocation {
    type Err = Box<Error>;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        s.parse().map(FromLocation::Integer)
                 .or_else(|_| s.parse().map(FromLocation::Register))
    }
}

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Register {
    A,
    B,
    C,
    D,
}

impl FromStr for Register {
    type Err = Box<Error>;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        use Register::*;
        match s {
            "a" => Ok(A),
            "b" => Ok(B),
            "c" => Ok(C),
            "d" => Ok(D),
            other => Err(format!("Unknown register {}", other).into()),
        }
    }
}

pub struct Machine {
    registers: [i32; 4],
}

impl Machine {
    pub fn new() -> Machine {
        Machine {
            registers: [0, 0, 1, 0],
        }
    }

    fn execute(&mut self, instructions: &[Instruction]) {
        let mut program_counter = 0;
        while let Some(&instruction) = instructions.get(program_counter) {
            // println!("registers = {:?}", self.registers);
            // println!("executing instruction at {}: {:?}", program_counter, instruction);

            let mut jump = 0;
            match instruction {
                Instruction::Copy(
                    FromLocation::Integer(i),
                    to
                ) => {
                    let index = Machine::register_index(to);
                    self.registers[index] = i;
                },
                Instruction::Copy(FromLocation::Register(from), to) => {
                    let index_from = Machine::register_index(from);
                    let index_to = Machine::register_index(to);
                    self.registers[index_to] = self.registers[index_from];
                },
                Instruction::Increment(register) => {
                    let index = Machine::register_index(register);
                    self.registers[index] += 1;
                },
                Instruction::Decrement(register) => {
                    let index = Machine::register_index(register);
                    self.registers[index] -= 1;
                },
                Instruction::JumpNonZero(FromLocation::Integer(i), offset) => {
                    if i != 0 {
                        jump = offset - 1;
                    }
                },
                Instruction::JumpNonZero(FromLocation::Register(register), offset) => {
                    let index = Machine::register_index(register);
                    if self.registers[index] != 0 {
                        jump = offset - 1;
                    }
                },
            }
            let next_pc = program_counter as i32 + 1 + jump;
            if next_pc < 0 {
                panic!("Jumped too far back!")
            }
            program_counter = next_pc as usize;
        }
    }

    fn register_index(register: Register) -> usize {
        use Register::*;
        match register {
            A => 0,
            B => 1,
            C => 2,
            D => 3,
        }
    }

    fn value_of(&self, register: Register) -> i32 {
        let index = Machine::register_index(register);
        self.registers[index]
    }
}

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

    #[test]
    fn parse_cpy() {
        let instr: Instruction = "cpy 41 a".parse().expect("couldn't parse");
        assert_eq!(
            instr,
            Instruction::Copy(FromLocation::Integer(41), Register::A)
        );
    }

    #[test]
    fn parse_cpy_register() {
        let instr: Instruction = "cpy a c".parse().expect("couldn't parse");
        assert_eq!(
            instr,
            Instruction::Copy(
                FromLocation::Register(Register::A),
                Register::C,
            )
        );
    }

    #[test]
    fn parse_inc() {
        let instr: Instruction = "inc a".parse().expect("couldn't parse");
        assert_eq!(
            instr,
            Instruction::Increment(Register::A)
        );
    }

    #[test]
    fn parse_dec() {
        let instr: Instruction = "dec a".parse().expect("couldn't parse");
        assert_eq!(
            instr,
            Instruction::Decrement(Register::A)
        );
    }

    #[test]
    fn parse_jnz() {
        let instr: Instruction = "jnz a 2".parse().expect("couldn't parse");
        assert_eq!(
            instr,
            Instruction::JumpNonZero(FromLocation::Register(Register::A), 2)
        );
    }

    #[test]
    fn execute_copy() {
        let mut machine = Machine::new();
        machine.execute(
            &[
                Instruction::Copy(
                    FromLocation::Integer(5),
                    Register::A,
                ),
            ]
        );

        assert_eq!(machine.value_of(Register::A), 5);
    }

    #[test]
    fn execute_copy_register() {
        let mut machine = Machine::new();
        machine.execute(
            &[
                Instruction::Copy(
                    FromLocation::Integer(5),
                    Register::A,
                ),
                Instruction::Copy(
                    FromLocation::Register(Register::A),
                    Register::B,
                ),
            ]
        );

        assert_eq!(machine.value_of(Register::B), 5);
    }

    #[test]
    fn execute_increment() {
        let mut machine = Machine::new();
        machine.execute(
            &[
                Instruction::Increment(
                    Register::C,
                ),
            ]
        );

        assert_eq!(machine.value_of(Register::C), 1);
    }

    #[test]
    fn execute_decrement() {
        let mut machine = Machine::new();
        machine.execute(
            &[
                Instruction::Decrement(
                    Register::D,
                ),
            ]
        );

        assert_eq!(machine.value_of(Register::D), -1);
    }

    #[test]
    fn execute_jnz_when_zero() {
        let mut machine = Machine::new();
        machine.execute(
            &[
                // D should be 0, jump doesn't happen
                Instruction::JumpNonZero(FromLocation::Register(Register::D), 2),
                Instruction::Increment(
                    Register::A,
                ),
            ]
        );

        assert_eq!(machine.value_of(Register::A), 1);
    }

    #[test]
    fn execute_jnz_when_non_zero() {
        let mut machine = Machine::new();
        machine.execute(
            &[
                Instruction::Decrement(
                    Register::D,
                ),
                Instruction::JumpNonZero(FromLocation::Register(Register::D), 2),
                Instruction::Increment(
                    Register::A,
                ),
            ]
        );

        assert_eq!(machine.value_of(Register::A), 0);
    }

    #[test]
    fn sample() {
        let input = "cpy 41 a
            inc a
            inc a
            dec a
            jnz a 2
            dec a";
        assert_eq!(puzzle(input), 42);
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// 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/.

#[macro_use]
extern crate log;
extern crate fern;
extern crate time;
extern crate libc;
extern crate rand;
extern crate simple_stream;
extern crate num_cpus;
//#[cfg(target_os = "linux")]
extern crate epoll;
extern crate rustc_serialize;

//#[cfg(target_os = "linux")]
pub mod server;
pub mod connectionloop;
pub mod epollloop;
pub mod resources;
pub mod socket;
pub mod stats;
pub mod types;
pub mod workerthread;
pub mod ipc;

/// Initializes all the global things
#[no_mangle]
pub extern "C" fn hydrogen_init() {
    println!("Initializing hydrogen logger...");

    let _ = fern::init_global_logger(fern::DispatchConfig {
        format: Box::new(|msg: &str, level: &log::LogLevel, _location: &log::LogLocation| {
            format!("[{}][{}] {}", time::now().strftime("%Y-%m-%d][%H:%M:%S").unwrap(), level, msg)
        }),
        output: vec![fern::OutputConfig::stdout(), fern::OutputConfig::file("/var/log/hydrogen.log")],
        level: log::LogLevelFilter::Trace,
    }, log::LogLevelFilter::Trace);

    trace!("Hydrogen initialized :)");
    info!("Log file: /var/log/hydrogen.log");
}
<commit_msg>GOing home<commit_after>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// 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/.

#[macro_use]
extern crate log;
extern crate fern;
extern crate time;
extern crate libc;
extern crate rand;
extern crate simple_stream;
extern crate num_cpus;
//#[cfg(target_os = "linux")]
extern crate epoll;
extern crate rustc_serialize;

//#[cfg(target_os = "linux")]
pub mod server;
pub mod connectionloop;
pub mod epollloop;
pub mod resources;
pub mod socket;
pub mod stats;
pub mod types;
pub mod workerthread;
pub mod ipc;

#[no_mangle]
pub extern "C" fn hydrogen_init() {
    let _ = fern::init_global_logger(fern::DispatchConfig {
        format: Box::new(|msg: &str, level: &log::LogLevel, _location: &log::LogLocation| {
            format!("[{}][{}] {}", time::now().strftime("%Y-%m-%d][%H:%M:%S").unwrap(), level, msg)
        }),
        output: vec![fern::OutputConfig::stdout(), fern::OutputConfig::file("/var/log/hydrogen.log")],
        level: log::LogLevelFilter::Trace,
    }, log::LogLevelFilter::Trace);
    info!("Initialized. \nLog file: /var/log/hydrogen.log");
}
<|endoftext|>


<commit_before>#![cfg_attr(feature = "nightly", feature(drain))]
#![deny(missing_docs)]
#![cfg_attr(test, deny(dead_code))]
#![cfg_attr(not(test), allow(dead_code))]

//! Small wrapper for `thread::spawn` that optionally auto-joins threads

use std::thread;
#[cfg(not(feature = "nightly"))] use std::mem;

/// A wrapper for `thread::spawn` that optionally auto-joins threads.
pub struct Spawner {
    threads: Vec<thread::JoinHandle<()>>
}

impl Spawner {
    /// Create a new Spawner object
    pub fn new() -> Spawner { Spawner { threads: vec![] } }

    /// Spawn a thread that will be auto-joined when the Spawner is dropped
    ///
    /// The thread function should be a move closure returning ()
    pub fn spawn_collected<F>(&mut self, f: F)
    where F: FnOnce(), F: Send + 'static
    {
        self.threads.push(thread::spawn(f));
    }

    /// Spawn a thread that won't be auto-joined
    ///
    /// The thread function should be a move closure
    pub fn spawn<F, T>(&mut self, f: F) -> thread::JoinHandle<T>
    where F: FnOnce() -> T, F: Send + 'static, T: Send + 'static
    {
        thread::spawn(f)
    }
}

impl Drop for Spawner {
    #[cfg(feature = "nightly")]
    fn drop(&mut self) {
        self.threads.drain(..)
            .map(thread::JoinHandle::join)
            .collect::<Result<Vec<_>,_>>()
            .unwrap();
    }

    #[cfg(not(feature = "nightly"))]
    fn drop(&mut self) {
        mem::replace(&mut self.threads, vec![])
            .into_iter()
            .map(thread::JoinHandle::join)
            .collect::<Result<Vec<_>,_>>()
            .unwrap();
    }
}

#[cfg(test)]
#[test]
fn spawn_some_threads() {
    use std::sync::atomic::{AtomicUsize, ATOMIC_USIZE_INIT, Ordering};
    use std::thread::sleep;
    use std::time::Duration;

    static ACTIVE: AtomicUsize = ATOMIC_USIZE_INIT;

    {
        let mut spawner = Spawner::new();

        assert!(
            spawner.spawn(move || {
                            ACTIVE.fetch_add(1, Ordering::SeqCst);
                            sleep(Duration::from_millis(100));
                            ACTIVE.fetch_sub(1, Ordering::SeqCst);
                         })
                   .join()
                   .is_ok()
        );

        assert_eq!(
            ACTIVE.load(Ordering::SeqCst),
            0
        );

        for _ in 1..10 {
            spawner.spawn_collected(move || {
                                      ACTIVE.fetch_add(1, Ordering::SeqCst);
                                      sleep(Duration::from_millis(100));
                                      ACTIVE.fetch_sub(1, Ordering::SeqCst);
                                   });
        }

        assert!(
            ACTIVE.load(Ordering::SeqCst) > 0
        );
    }

    assert_eq!(
        ACTIVE.load(Ordering::SeqCst),
        0
    );
}

<commit_msg>don't allocate in drop method<commit_after>#![cfg_attr(feature = "nightly", feature(drain))]
#![deny(missing_docs)]
#![cfg_attr(test, deny(dead_code))]
#![cfg_attr(not(test), allow(dead_code))]

//! Small wrapper for `thread::spawn` that optionally auto-joins threads

use std::thread;
#[cfg(not(feature = "nightly"))] use std::mem;

/// A wrapper for `thread::spawn` that optionally auto-joins threads.
pub struct Spawner {
    threads: Vec<thread::JoinHandle<()>>
}

impl Spawner {
    /// Create a new Spawner object
    pub fn new() -> Spawner { Spawner { threads: vec![] } }

    /// Spawn a thread that will be auto-joined when the Spawner is dropped
    ///
    /// The thread function should be a move closure returning ()
    pub fn spawn_collected<F>(&mut self, f: F)
    where F: FnOnce(), F: Send + 'static
    {
        self.threads.push(thread::spawn(f));
    }

    /// Spawn a thread that won't be auto-joined
    ///
    /// The thread function should be a move closure
    pub fn spawn<F, T>(&mut self, f: F) -> thread::JoinHandle<T>
    where F: FnOnce() -> T, F: Send + 'static, T: Send + 'static
    {
        thread::spawn(f)
    }
}

impl Drop for Spawner {
    #[cfg(feature = "nightly")]
    fn drop(&mut self) {
        assert!(self.threads.drain(..)
                    .map(thread::JoinHandle::join)
                    .find(Result::is_err)
                    .is_none());
    }

    #[cfg(not(feature = "nightly"))]
    fn drop(&mut self) {
        assert!(mem::replace(&mut self.threads, vec![])
                    .into_iter()
                    .map(thread::JoinHandle::join)
                    .find(Result::is_err)
                    .is_none());
    }
}

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

    #[test]
    fn spawn_some_threads() {
        use std::sync::atomic::{AtomicUsize, ATOMIC_USIZE_INIT, Ordering};
        use std::thread::sleep;
        use std::time::Duration;

        static ACTIVE: AtomicUsize = ATOMIC_USIZE_INIT;

        {
            let mut spawner = Spawner::new();

            // an manually joined thread
            assert!(
                spawner.spawn(move || {
                                ACTIVE.fetch_add(1, Ordering::SeqCst);
                                sleep(Duration::from_millis(100));
                                ACTIVE.fetch_sub(1, Ordering::SeqCst);
                             })
                       .join()
                       .is_ok()
            );

            // make sure the manual thread finished
            assert_eq!(
                ACTIVE.load(Ordering::SeqCst),
                0
            );

            // some collected threads
            for _ in 1..10 {
                spawner.spawn_collected(move || {
                                          ACTIVE.fetch_add(1, Ordering::SeqCst);
                                          sleep(Duration::from_millis(100));
                                          ACTIVE.fetch_sub(1, Ordering::SeqCst);
                                       });
            }

            // make sure at least some have started
            assert!(
                ACTIVE.load(Ordering::SeqCst) > 0
            );

            // collected threads implicitly joined here
        }

        // make sure they all finished
        assert_eq!(
            ACTIVE.load(Ordering::SeqCst),
            0
        );
    }

    #[test]
    #[should_panic]
    fn failing_thread() {
        {
            let mut spawner = Spawner::new();
            spawner.spawn_collected(move || panic!());
        }
    }
}

<|endoftext|>