<commit_before>use crate::colors::ColorType;
use crate::deflate;
use crate::error::PngError;
use crate::filters::*;
use crate::headers::*;
use crate::interlace::{deinterlace_image, interlace_image};
use byteorder::{BigEndian, WriteBytesExt};
use crc::crc32;
use indexmap::IndexMap;
use rgb::ComponentSlice;
use rgb::RGBA8;
use std::fs::File;
use std::io::{BufReader, Read};
use std::iter::Iterator;
use std::path::Path;
use std::sync::Arc;

pub(crate) const STD_COMPRESSION: u8 = 6;
/// Must use normal compression, as faster ones (Huffman/RLE-only) are not representative
pub(crate) const STD_STRATEGY: u8 = 0;
/// OK to use a bit smalller window for evaluation
pub(crate) const STD_WINDOW: u8 = 13;
pub(crate) const STD_FILTERS: [u8; 2] = [0, 5];

pub(crate) mod scan_lines;

use self::scan_lines::{ScanLines, ScanLinesMut};

#[derive(Debug, Clone)]
pub struct PngImage {
    /// The headers stored in the IHDR chunk
    pub ihdr: IhdrData,
    /// The uncompressed, optionally filtered data from the IDAT chunk
    pub data: Vec<u8>,
    /// The palette containing colors used in an Indexed image
    /// Contains 3 bytes per color (R+G+B), up to 768
    pub palette: Option<Vec<RGBA8>>,
    /// The pixel value that should be rendered as transparent
    pub transparency_pixel: Option<Vec<u8>>,
    /// All non-critical headers from the PNG are stored here
    pub aux_headers: IndexMap<[u8; 4], Vec<u8>>,
}

/// Contains all data relevant to a PNG image
#[derive(Debug, Clone)]
pub struct PngData {
    /// Uncompressed image data
    pub raw: Arc<PngImage>,
    /// The filtered and compressed data of the IDAT chunk
    pub idat_data: Vec<u8>,
}

type PaletteWithTrns = (Option<Vec<RGBA8>>, Option<Vec<u8>>);

impl PngData {
    /// Create a new `PngData` struct by opening a file
    #[inline]
    pub fn new(filepath: &Path, fix_errors: bool) -> Result<Self, PngError> {
        let byte_data = Self::read_file(filepath)?;

        Self::from_slice(&byte_data, fix_errors)
    }

    pub fn read_file(filepath: &Path) -> Result<Vec<u8>, PngError> {
        let file = match File::open(filepath) {
            Ok(f) => f,
            Err(_) => return Err(PngError::new("Failed to open file for reading")),
        };
        let file_len = file.metadata().map(|m| m.len() as usize).unwrap_or(0);
        let mut reader = BufReader::new(file);
        // Check file for PNG header
        let mut header = [0; 8];
        if reader.read_exact(&mut header).is_err() {
            return Err(PngError::new("Not a PNG file: too small"));
        }
        if !file_header_is_valid(&header) {
            return Err(PngError::new("Invalid PNG header detected"));
        }
        // Read raw png data into memory
        let mut byte_data: Vec<u8> = Vec::with_capacity(file_len);
        byte_data.extend_from_slice(&header);
        match reader.read_to_end(&mut byte_data) {
            Ok(_) => (),
            Err(_) => return Err(PngError::new("Failed to read from file")),
        }
        Ok(byte_data)
    }

    /// Create a new `PngData` struct by reading a slice
    pub fn from_slice(byte_data: &[u8], fix_errors: bool) -> Result<Self, PngError> {
        let mut byte_offset: usize = 0;
        // Test that png header is valid
        let header = byte_data.get(0..8).ok_or(PngError::TruncatedData)?;
        if !file_header_is_valid(header) {
            return Err(PngError::NotPNG);
        }
        byte_offset += 8;
        // Read the data headers
        let mut aux_headers: IndexMap<[u8; 4], Vec<u8>> = IndexMap::new();
        let mut idat_headers: Vec<u8> = Vec::new();
        while let Some(header) = parse_next_header(byte_data, &mut byte_offset, fix_errors)? {
            match &header.name {
                b"IDAT" => idat_headers.extend_from_slice(header.data),
                b"acTL" => return Err(PngError::APNGNotSupported),
                _ => {
                    aux_headers.insert(header.name, header.data.to_owned());
                }
            }
        }
        // Parse the headers into our PngData
        if idat_headers.is_empty() {
            return Err(PngError::ChunkMissing("IDAT"));
        }
        let ihdr = match aux_headers.remove(b"IHDR") {
            Some(ihdr) => ihdr,
            None => return Err(PngError::ChunkMissing("IHDR")),
        };
        let ihdr_header = parse_ihdr_header(&ihdr)?;
        let raw_data = deflate::inflate(idat_headers.as_ref())?;

        // Reject files with incorrect width/height or truncated data
        if raw_data.len() != ihdr_header.raw_data_size() {
            return Err(PngError::TruncatedData);
        }

        let (palette, transparency_pixel) = Self::palette_to_rgba(
            ihdr_header.color_type,
            aux_headers.remove(b"PLTE"),
            aux_headers.remove(b"tRNS"),
        )?;

        let mut raw = PngImage {
            ihdr: ihdr_header,
            data: raw_data,
            palette,
            transparency_pixel,
            aux_headers,
        };
        raw.data = raw.unfilter_image();
        // Return the PngData
        Ok(Self {
            idat_data: idat_headers,
            raw: Arc::new(raw),
        })
    }

    /// Handle transparency header
    fn palette_to_rgba(
        color_type: ColorType,
        palette_data: Option<Vec<u8>>,
        trns_data: Option<Vec<u8>>,
    ) -> Result<PaletteWithTrns, PngError> {
        if color_type == ColorType::Indexed {
            let palette_data =
                palette_data.ok_or_else(|| PngError::new("no palette in indexed image"))?;
            let mut palette: Vec<_> = palette_data
                .chunks(3)
                .map(|color| RGBA8::new(color[0], color[1], color[2], 255))
                .collect();

            if let Some(trns_data) = trns_data {
                for (color, trns) in palette.iter_mut().zip(trns_data) {
                    color.a = trns;
                }
            }
            Ok((Some(palette), None))
        } else {
            Ok((None, trns_data))
        }
    }

    /// Format the `PngData` struct into a valid PNG bytestream
    pub fn output(&self) -> Vec<u8> {
        // PNG header
        let mut output = vec![0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
        // IHDR
        let mut ihdr_data = Vec::with_capacity(13);
        let _ = ihdr_data.write_u32::<BigEndian>(self.raw.ihdr.width);
        let _ = ihdr_data.write_u32::<BigEndian>(self.raw.ihdr.height);
        let _ = ihdr_data.write_u8(self.raw.ihdr.bit_depth.as_u8());
        let _ = ihdr_data.write_u8(self.raw.ihdr.color_type.png_header_code());
        let _ = ihdr_data.write_u8(0); // Compression -- deflate
        let _ = ihdr_data.write_u8(0); // Filter method -- 5-way adaptive filtering
        let _ = ihdr_data.write_u8(self.raw.ihdr.interlaced);
        write_png_block(b"IHDR", &ihdr_data, &mut output);
        // Ancillary headers
        for (key, header) in self
            .raw
            .aux_headers
            .iter()
            .filter(|&(key, _)| !(key == b"bKGD" || key == b"hIST" || key == b"tRNS"))
        {
            write_png_block(key, header, &mut output);
        }
        // Palette
        if let Some(ref palette) = self.raw.palette {
            let mut palette_data = Vec::with_capacity(palette.len() * 3);
            let max_palette_size = 1 << (self.raw.ihdr.bit_depth.as_u8() as usize);
            for px in palette.iter().take(max_palette_size) {
                palette_data.extend_from_slice(px.rgb().as_slice());
            }
            write_png_block(b"PLTE", &palette_data, &mut output);
            let num_transparent =
                palette
                    .iter()
                    .take(max_palette_size)
                    .enumerate()
                    .fold(
                        0,
                        |prev, (index, px)| {
                            if px.a == 255 {
                                prev
                            } else {
                                index + 1
                            }
                        },
                    );
            if num_transparent > 0 {
                let trns_data: Vec<_> = palette[0..num_transparent].iter().map(|px| px.a).collect();
                write_png_block(b"tRNS", &trns_data, &mut output);
            }
        } else if let Some(ref transparency_pixel) = self.raw.transparency_pixel {
            // Transparency pixel
            write_png_block(b"tRNS", transparency_pixel, &mut output);
        }
        // Special ancillary headers that need to come after PLTE but before IDAT
        for (key, header) in self
            .raw
            .aux_headers
            .iter()
            .filter(|&(key, _)| key == b"bKGD" || key == b"hIST" || key == b"tRNS")
        {
            write_png_block(key, header, &mut output);
        }
        // IDAT data
        write_png_block(b"IDAT", &self.idat_data, &mut output);
        // Stream end
        write_png_block(b"IEND", &[], &mut output);

        output
    }
}

impl PngImage {
    /// Convert the image to the specified interlacing type
    /// Returns true if the interlacing was changed, false otherwise
    /// The `interlace` parameter specifies the *new* interlacing mode
    /// Assumes that the data has already been de-filtered
    #[inline]
    #[must_use]
    pub fn change_interlacing(&self, interlace: u8) -> Option<PngImage> {
        if interlace == self.ihdr.interlaced {
            return None;
        }

        Some(if interlace == 1 {
            // Convert progressive to interlaced data
            interlace_image(self)
        } else {
            // Convert interlaced to progressive data
            deinterlace_image(self)
        })
    }

    /// Return the number of channels in the image, based on color type
    #[inline]
    pub fn channels_per_pixel(&self) -> u8 {
        self.ihdr.color_type.channels_per_pixel()
    }

    /// Return an iterator over the scanlines of the image
    #[inline]
    pub fn scan_lines(&self) -> ScanLines<'_> {
        ScanLines::new(self)
    }

    /// Return an iterator over the scanlines of the image
    #[inline]
    pub fn scan_lines_mut(&mut self) -> ScanLinesMut<'_> {
        ScanLinesMut::new(self)
    }

    /// Reverse all filters applied on the image, returning an unfiltered IDAT bytestream
    fn unfilter_image(&self) -> Vec<u8> {
        let mut unfiltered = Vec::with_capacity(self.data.len());
        let bpp = ((self.ihdr.bit_depth.as_u8() * self.channels_per_pixel() + 7) / 8) as usize;
        let mut last_line: Vec<u8> = Vec::new();
        let mut last_pass = 1;
        let mut unfiltered_buf = Vec::new();
        for line in self.scan_lines() {
            if let Some(pass) = line.pass {
                if pass != last_pass {
                    last_line.clear();
                    last_pass = pass;
                }
            }
            last_line.resize(line.data.len(), 0);
            unfilter_line(
                line.filter,
                bpp,
                &line.data,
                &last_line,
                &mut unfiltered_buf,
            );
            unfiltered.push(0);
            unfiltered.extend_from_slice(&unfiltered_buf);
            std::mem::swap(&mut last_line, &mut unfiltered_buf);
            unfiltered_buf.clear();
        }
        unfiltered
    }

    /// Apply the specified filter type to all rows in the image
    /// 0: None
    /// 1: Sub
    /// 2: Up
    /// 3: Average
    /// 4: Paeth
    /// 5: All (heuristically pick the best filter for each line)
    pub fn filter_image(&self, filter: u8) -> Vec<u8> {
        let mut filtered = Vec::with_capacity(self.data.len());
        let bpp = ((self.ihdr.bit_depth.as_u8() * self.channels_per_pixel() + 7) / 8) as usize;
        let mut last_line: &[u8] = &[];
        let mut last_pass: Option<u8> = None;
        let mut f_buf = Vec::new();
        for line in self.scan_lines() {
            f_buf.clear();
            match filter {
                0 | 1 | 2 | 3 | 4 => {
                    let filter = if last_pass == line.pass || filter <= 1 {
                        filter
                    } else {
                        0
                    };
                    filtered.push(filter);
                    filter_line(filter, bpp, &line.data, last_line, &mut f_buf);
                    filtered.extend_from_slice(&f_buf);
                }
                5 => {
                    // Heuristically guess best filter per line
                    // Uses MSAD algorithm mentioned in libpng reference docs
                    // http://www.libpng.org/pub/png/book/chapter09.html
                    let mut best_filter = 0;
                    let mut best_line = Vec::new();
                    let mut best_size = std::u64::MAX;

                    // Avoid vertical filtering on first line of each interlacing pass
                    for filter in if last_pass == line.pass { 0..5 } else { 0..2 } {
                        filter_line(filter, bpp, &line.data, last_line, &mut f_buf);
                        let size = f_buf.iter().fold(0_u64, |acc, &x| {
                            let signed = x as i8;
                            acc + i16::from(signed).abs() as u64
                        });
                        if size < best_size {
                            best_size = size;
                            best_filter = filter;
                            std::mem::swap(&mut best_line, &mut f_buf);
                        }
                        f_buf.clear() //discard buffer, and start again
                    }
                    filtered.push(best_filter);
                    filtered.extend_from_slice(&best_line);
                }
                _ => unreachable!(),
            }
            last_line = line.data;
            last_pass = line.pass;
        }
        filtered
    }
}
fn write_png_block(key: &[u8], header: &[u8], output: &mut Vec<u8>) {
    let mut header_data = Vec::with_capacity(header.len() + 4);
    header_data.extend_from_slice(key);
    header_data.extend_from_slice(header);
    output.reserve(header_data.len() + 8);
    let _ = output.write_u32::<BigEndian>(header_data.len() as u32 - 4);
    let crc = crc32::checksum_ieee(&header_data);
    output.append(&mut header_data);
    let _ = output.write_u32::<BigEndian>(crc);
}
<commit_msg>Make sure passes clear last line<commit_after>use crate::colors::ColorType;
use crate::deflate;
use crate::error::PngError;
use crate::filters::*;
use crate::headers::*;
use crate::interlace::{deinterlace_image, interlace_image};
use byteorder::{BigEndian, WriteBytesExt};
use crc::crc32;
use indexmap::IndexMap;
use rgb::ComponentSlice;
use rgb::RGBA8;
use std::fs::File;
use std::io::{BufReader, Read};
use std::iter::Iterator;
use std::path::Path;
use std::sync::Arc;

pub(crate) const STD_COMPRESSION: u8 = 6;
/// Must use normal compression, as faster ones (Huffman/RLE-only) are not representative
pub(crate) const STD_STRATEGY: u8 = 0;
/// OK to use a bit smalller window for evaluation
pub(crate) const STD_WINDOW: u8 = 13;
pub(crate) const STD_FILTERS: [u8; 2] = [0, 5];

pub(crate) mod scan_lines;

use self::scan_lines::{ScanLines, ScanLinesMut};

#[derive(Debug, Clone)]
pub struct PngImage {
    /// The headers stored in the IHDR chunk
    pub ihdr: IhdrData,
    /// The uncompressed, optionally filtered data from the IDAT chunk
    pub data: Vec<u8>,
    /// The palette containing colors used in an Indexed image
    /// Contains 3 bytes per color (R+G+B), up to 768
    pub palette: Option<Vec<RGBA8>>,
    /// The pixel value that should be rendered as transparent
    pub transparency_pixel: Option<Vec<u8>>,
    /// All non-critical headers from the PNG are stored here
    pub aux_headers: IndexMap<[u8; 4], Vec<u8>>,
}

/// Contains all data relevant to a PNG image
#[derive(Debug, Clone)]
pub struct PngData {
    /// Uncompressed image data
    pub raw: Arc<PngImage>,
    /// The filtered and compressed data of the IDAT chunk
    pub idat_data: Vec<u8>,
}

type PaletteWithTrns = (Option<Vec<RGBA8>>, Option<Vec<u8>>);

impl PngData {
    /// Create a new `PngData` struct by opening a file
    #[inline]
    pub fn new(filepath: &Path, fix_errors: bool) -> Result<Self, PngError> {
        let byte_data = Self::read_file(filepath)?;

        Self::from_slice(&byte_data, fix_errors)
    }

    pub fn read_file(filepath: &Path) -> Result<Vec<u8>, PngError> {
        let file = match File::open(filepath) {
            Ok(f) => f,
            Err(_) => return Err(PngError::new("Failed to open file for reading")),
        };
        let file_len = file.metadata().map(|m| m.len() as usize).unwrap_or(0);
        let mut reader = BufReader::new(file);
        // Check file for PNG header
        let mut header = [0; 8];
        if reader.read_exact(&mut header).is_err() {
            return Err(PngError::new("Not a PNG file: too small"));
        }
        if !file_header_is_valid(&header) {
            return Err(PngError::new("Invalid PNG header detected"));
        }
        // Read raw png data into memory
        let mut byte_data: Vec<u8> = Vec::with_capacity(file_len);
        byte_data.extend_from_slice(&header);
        match reader.read_to_end(&mut byte_data) {
            Ok(_) => (),
            Err(_) => return Err(PngError::new("Failed to read from file")),
        }
        Ok(byte_data)
    }

    /// Create a new `PngData` struct by reading a slice
    pub fn from_slice(byte_data: &[u8], fix_errors: bool) -> Result<Self, PngError> {
        let mut byte_offset: usize = 0;
        // Test that png header is valid
        let header = byte_data.get(0..8).ok_or(PngError::TruncatedData)?;
        if !file_header_is_valid(header) {
            return Err(PngError::NotPNG);
        }
        byte_offset += 8;
        // Read the data headers
        let mut aux_headers: IndexMap<[u8; 4], Vec<u8>> = IndexMap::new();
        let mut idat_headers: Vec<u8> = Vec::new();
        while let Some(header) = parse_next_header(byte_data, &mut byte_offset, fix_errors)? {
            match &header.name {
                b"IDAT" => idat_headers.extend_from_slice(header.data),
                b"acTL" => return Err(PngError::APNGNotSupported),
                _ => {
                    aux_headers.insert(header.name, header.data.to_owned());
                }
            }
        }
        // Parse the headers into our PngData
        if idat_headers.is_empty() {
            return Err(PngError::ChunkMissing("IDAT"));
        }
        let ihdr = match aux_headers.remove(b"IHDR") {
            Some(ihdr) => ihdr,
            None => return Err(PngError::ChunkMissing("IHDR")),
        };
        let ihdr_header = parse_ihdr_header(&ihdr)?;
        let raw_data = deflate::inflate(idat_headers.as_ref())?;

        // Reject files with incorrect width/height or truncated data
        if raw_data.len() != ihdr_header.raw_data_size() {
            return Err(PngError::TruncatedData);
        }

        let (palette, transparency_pixel) = Self::palette_to_rgba(
            ihdr_header.color_type,
            aux_headers.remove(b"PLTE"),
            aux_headers.remove(b"tRNS"),
        )?;

        let mut raw = PngImage {
            ihdr: ihdr_header,
            data: raw_data,
            palette,
            transparency_pixel,
            aux_headers,
        };
        raw.data = raw.unfilter_image();
        // Return the PngData
        Ok(Self {
            idat_data: idat_headers,
            raw: Arc::new(raw),
        })
    }

    /// Handle transparency header
    fn palette_to_rgba(
        color_type: ColorType,
        palette_data: Option<Vec<u8>>,
        trns_data: Option<Vec<u8>>,
    ) -> Result<PaletteWithTrns, PngError> {
        if color_type == ColorType::Indexed {
            let palette_data =
                palette_data.ok_or_else(|| PngError::new("no palette in indexed image"))?;
            let mut palette: Vec<_> = palette_data
                .chunks(3)
                .map(|color| RGBA8::new(color[0], color[1], color[2], 255))
                .collect();

            if let Some(trns_data) = trns_data {
                for (color, trns) in palette.iter_mut().zip(trns_data) {
                    color.a = trns;
                }
            }
            Ok((Some(palette), None))
        } else {
            Ok((None, trns_data))
        }
    }

    /// Format the `PngData` struct into a valid PNG bytestream
    pub fn output(&self) -> Vec<u8> {
        // PNG header
        let mut output = vec![0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
        // IHDR
        let mut ihdr_data = Vec::with_capacity(13);
        let _ = ihdr_data.write_u32::<BigEndian>(self.raw.ihdr.width);
        let _ = ihdr_data.write_u32::<BigEndian>(self.raw.ihdr.height);
        let _ = ihdr_data.write_u8(self.raw.ihdr.bit_depth.as_u8());
        let _ = ihdr_data.write_u8(self.raw.ihdr.color_type.png_header_code());
        let _ = ihdr_data.write_u8(0); // Compression -- deflate
        let _ = ihdr_data.write_u8(0); // Filter method -- 5-way adaptive filtering
        let _ = ihdr_data.write_u8(self.raw.ihdr.interlaced);
        write_png_block(b"IHDR", &ihdr_data, &mut output);
        // Ancillary headers
        for (key, header) in self
            .raw
            .aux_headers
            .iter()
            .filter(|&(key, _)| !(key == b"bKGD" || key == b"hIST" || key == b"tRNS"))
        {
            write_png_block(key, header, &mut output);
        }
        // Palette
        if let Some(ref palette) = self.raw.palette {
            let mut palette_data = Vec::with_capacity(palette.len() * 3);
            let max_palette_size = 1 << (self.raw.ihdr.bit_depth.as_u8() as usize);
            for px in palette.iter().take(max_palette_size) {
                palette_data.extend_from_slice(px.rgb().as_slice());
            }
            write_png_block(b"PLTE", &palette_data, &mut output);
            let num_transparent =
                palette
                    .iter()
                    .take(max_palette_size)
                    .enumerate()
                    .fold(
                        0,
                        |prev, (index, px)| {
                            if px.a == 255 {
                                prev
                            } else {
                                index + 1
                            }
                        },
                    );
            if num_transparent > 0 {
                let trns_data: Vec<_> = palette[0..num_transparent].iter().map(|px| px.a).collect();
                write_png_block(b"tRNS", &trns_data, &mut output);
            }
        } else if let Some(ref transparency_pixel) = self.raw.transparency_pixel {
            // Transparency pixel
            write_png_block(b"tRNS", transparency_pixel, &mut output);
        }
        // Special ancillary headers that need to come after PLTE but before IDAT
        for (key, header) in self
            .raw
            .aux_headers
            .iter()
            .filter(|&(key, _)| key == b"bKGD" || key == b"hIST" || key == b"tRNS")
        {
            write_png_block(key, header, &mut output);
        }
        // IDAT data
        write_png_block(b"IDAT", &self.idat_data, &mut output);
        // Stream end
        write_png_block(b"IEND", &[], &mut output);

        output
    }
}

impl PngImage {
    /// Convert the image to the specified interlacing type
    /// Returns true if the interlacing was changed, false otherwise
    /// The `interlace` parameter specifies the *new* interlacing mode
    /// Assumes that the data has already been de-filtered
    #[inline]
    #[must_use]
    pub fn change_interlacing(&self, interlace: u8) -> Option<PngImage> {
        if interlace == self.ihdr.interlaced {
            return None;
        }

        Some(if interlace == 1 {
            // Convert progressive to interlaced data
            interlace_image(self)
        } else {
            // Convert interlaced to progressive data
            deinterlace_image(self)
        })
    }

    /// Return the number of channels in the image, based on color type
    #[inline]
    pub fn channels_per_pixel(&self) -> u8 {
        self.ihdr.color_type.channels_per_pixel()
    }

    /// Return an iterator over the scanlines of the image
    #[inline]
    pub fn scan_lines(&self) -> ScanLines<'_> {
        ScanLines::new(self)
    }

    /// Return an iterator over the scanlines of the image
    #[inline]
    pub fn scan_lines_mut(&mut self) -> ScanLinesMut<'_> {
        ScanLinesMut::new(self)
    }

    /// Reverse all filters applied on the image, returning an unfiltered IDAT bytestream
    fn unfilter_image(&self) -> Vec<u8> {
        let mut unfiltered = Vec::with_capacity(self.data.len());
        let bpp = ((self.ihdr.bit_depth.as_u8() * self.channels_per_pixel() + 7) / 8) as usize;
        let mut last_line: Vec<u8> = Vec::new();
        let mut last_pass = None;
        let mut unfiltered_buf = Vec::new();
        for line in self.scan_lines() {
            if last_pass != line.pass {
                last_line.clear();
                last_pass = line.pass;
            }
            last_line.resize(line.data.len(), 0);
            unfilter_line(
                line.filter,
                bpp,
                &line.data,
                &last_line,
                &mut unfiltered_buf,
            );
            unfiltered.push(0);
            unfiltered.extend_from_slice(&unfiltered_buf);
            std::mem::swap(&mut last_line, &mut unfiltered_buf);
            unfiltered_buf.clear();
        }
        unfiltered
    }

    /// Apply the specified filter type to all rows in the image
    /// 0: None
    /// 1: Sub
    /// 2: Up
    /// 3: Average
    /// 4: Paeth
    /// 5: All (heuristically pick the best filter for each line)
    pub fn filter_image(&self, filter: u8) -> Vec<u8> {
        let mut filtered = Vec::with_capacity(self.data.len());
        let bpp = ((self.ihdr.bit_depth.as_u8() * self.channels_per_pixel() + 7) / 8) as usize;
        let mut last_line: &[u8] = &[];
        let mut last_pass: Option<u8> = None;
        let mut f_buf = Vec::new();
        for line in self.scan_lines() {
            f_buf.clear();
            if last_pass != line.pass {
                last_line = &[];
            }
            match filter {
                0 | 1 | 2 | 3 | 4 => {
                    let filter = if last_pass == line.pass || filter <= 1 {
                        filter
                    } else {
                        0
                    };
                    filtered.push(filter);
                    filter_line(filter, bpp, &line.data, last_line, &mut f_buf);
                    filtered.extend_from_slice(&f_buf);
                }
                5 => {
                    // Heuristically guess best filter per line
                    // Uses MSAD algorithm mentioned in libpng reference docs
                    // http://www.libpng.org/pub/png/book/chapter09.html
                    let mut best_filter = 0;
                    let mut best_line = Vec::new();
                    let mut best_size = std::u64::MAX;

                    // Avoid vertical filtering on first line of each interlacing pass
                    for filter in if last_pass == line.pass { 0..5 } else { 0..2 } {
                        filter_line(filter, bpp, &line.data, last_line, &mut f_buf);
                        let size = f_buf.iter().fold(0_u64, |acc, &x| {
                            let signed = x as i8;
                            acc + i16::from(signed).abs() as u64
                        });
                        if size < best_size {
                            best_size = size;
                            best_filter = filter;
                            std::mem::swap(&mut best_line, &mut f_buf);
                        }
                        f_buf.clear() //discard buffer, and start again
                    }
                    filtered.push(best_filter);
                    filtered.extend_from_slice(&best_line);
                }
                _ => unreachable!(),
            }
            last_line = line.data;
            last_pass = line.pass;
        }
        filtered
    }
}
fn write_png_block(key: &[u8], header: &[u8], output: &mut Vec<u8>) {
    let mut header_data = Vec::with_capacity(header.len() + 4);
    header_data.extend_from_slice(key);
    header_data.extend_from_slice(header);
    output.reserve(header_data.len() + 8);
    let _ = output.write_u32::<BigEndian>(header_data.len() as u32 - 4);
    let crc = crc32::checksum_ieee(&header_data);
    output.append(&mut header_data);
    let _ = output.write_u32::<BigEndian>(crc);
}
<|endoftext|>


<commit_before>//! A collection of commonly used items that we recommend importing for ease of use.

pub use app::{App, LoopMode};
pub use color;
pub use color::named::*;
pub use event::SimpleWindowEvent::*;
pub use event::{Event, Key};
pub use frame::Frame;
pub use math::{Point2, Point3, Vector2, Vector3, vec1, vec2, vec3, vec4, pt1, pt2, pt3};
pub use math::prelude::*;
pub use window::Id as WindowId;
<commit_msg>Remove color module from prelude to avoid ui conflicts<commit_after>//! A collection of commonly used items that we recommend importing for ease of use.

pub use app::{App, LoopMode};
pub use color::named::*;
pub use event::SimpleWindowEvent::*;
pub use event::{Event, Key};
pub use frame::Frame;
pub use math::{Point2, Point3, Vector2, Vector3, vec1, vec2, vec3, vec4, pt1, pt2, pt3};
pub use math::prelude::*;
pub use window::Id as WindowId;
<|endoftext|>


<commit_before>use std::io;
use std::io::Write;
use std::iter;

use byteorder::{LittleEndian, WriteBytesExt};
use encoding::all::UTF_8;
use encoding::types::{EncoderTrap, EncodingRef};

use consts;
use img::Image;

pub struct Printer<W: io::Write> {
    writer: io::BufWriter<W>,
    codec: EncodingRef,
    trap: EncoderTrap,
}

impl<W: io::Write> Printer<W> {
    pub fn new(writer: W, codec: Option<EncodingRef>, trap: Option<EncoderTrap>) -> Printer<W> {
        Printer {
            writer: io::BufWriter::new(writer),
            codec: codec.unwrap_or(UTF_8 as EncodingRef),
            trap: trap.unwrap_or(EncoderTrap::Replace),
        }
    }

    fn encode(&mut self, content: &str) -> Vec<u8> {
        self.codec.encode(content, self.trap).unwrap()
    }

    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.writer.write(buf)
    }

    pub fn write_u8(&mut self, n: u8) -> io::Result<usize> {
        self.write(vec![n].as_slice())
    }

    fn write_u16le(&mut self, n: u16) -> io::Result<usize> {
        let mut wtr = vec![];
        wtr.write_u16::<LittleEndian>(n).unwrap();
        self.write(wtr.as_slice())
    }

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

    pub fn hwinit(&mut self) -> io::Result<usize> {
        self.write(consts::HW_INIT)
    }

    pub fn hwselect(&mut self) -> io::Result<usize> {
        self.write(consts::HW_SELECT)
    }

    pub fn hwreset(&mut self) -> io::Result<usize> {
        self.write(consts::HW_RESET)
    }

    pub fn print(&mut self, content: &str) -> io::Result<usize> {
        // let rv = self.encode(content);
        let rv = self.encode(content);
        self.write(rv.as_slice())
    }

    pub fn println(&mut self, content: &str) -> io::Result<usize> {
        self.print(format!("{}{}", content, consts::EOL).as_ref())
    }

    pub fn text(&mut self, content: &str) -> io::Result<usize> {
        self.println(content)
    }

    pub fn line_space(&mut self, n: i32) -> io::Result<usize> {
        if n >= 0 {
            Ok(self.write(consts::LS_SET)? + self.write_u8(n as u8)?)
        } else {
            self.write(consts::LS_DEFAULT)
        }
    }

    pub fn feed(&mut self, n: usize) -> io::Result<usize> {
        let n = if n < 1 { 1 } else { n };
        self.write(
            iter::repeat(consts::EOL)
                .take(n)
                .collect::<String>()
                .as_ref(),
        )
    }

    pub fn control(&mut self, ctrl: &str) -> io::Result<usize> {
        let ctrl_upper = ctrl.to_uppercase();
        let ctrl_value = match ctrl_upper.as_ref() {
            "LF" => consts::CTL_LF,
            "FF" => consts::CTL_FF,
            "CR" => consts::CTL_CR,
            "HT" => consts::CTL_HT,
            "VT" => consts::CTL_VT,
            _ => panic!("Invalid control action: {}", ctrl),
        };
        self.write(ctrl_value)
    }

    pub fn align(&mut self, alignment: &str) -> io::Result<usize> {
        let align_upper = alignment.to_uppercase();
        let align_value = match align_upper.as_ref() {
            "LT" => consts::TXT_ALIGN_LT,
            "CT" => consts::TXT_ALIGN_CT,
            "RT" => consts::TXT_ALIGN_RT,
            _ => panic!("Invalid alignment: {}", alignment),
        };
        self.write(align_value)
    }

    pub fn font(&mut self, family: &str) -> io::Result<usize> {
        let family_upper = family.to_uppercase();
        let family_value = match family_upper.as_ref() {
            "A" => consts::TXT_FONT_A,
            "B" => consts::TXT_FONT_B,
            "C" => consts::TXT_FONT_C,
            _ => panic!("Invalid font family: {}", family),
        };
        self.write(family_value)
    }

    pub fn style(&mut self, kind: &str) -> io::Result<usize> {
        let kind_upper = kind.to_uppercase();
        match kind_upper.as_ref() {
            "B" => Ok(self.write(consts::TXT_UNDERL_OFF)? + self.write(consts::TXT_BOLD_ON)?),
            "U" => Ok(self.write(consts::TXT_BOLD_OFF)? + self.write(consts::TXT_UNDERL_ON)?),
            "U2" => Ok(self.write(consts::TXT_BOLD_OFF)? + self.write(consts::TXT_UNDERL2_ON)?),
            "BU" => Ok(self.write(consts::TXT_BOLD_ON)? + self.write(consts::TXT_UNDERL_ON)?),
            "BU2" => Ok(self.write(consts::TXT_BOLD_ON)? + self.write(consts::TXT_UNDERL2_ON)?),
            "NORMAL" | _ => {
                Ok(self.write(consts::TXT_BOLD_OFF)? + self.write(consts::TXT_UNDERL_OFF)?)
            }
        }
    }

    pub fn size(&mut self, width: usize, height: usize) -> io::Result<usize> {
        let mut n = self.write(consts::TXT_NORMAL)?;
        if width == 2 {
            n += self.write(consts::TXT_2WIDTH)?;
        }
        if height == 2 {
            n += self.write(consts::TXT_2HEIGHT)?;
        }
        Ok(n)
    }

    pub fn hardware(&mut self, hw: &str) -> io::Result<usize> {
        let value = match hw {
            "INIT" => consts::HW_INIT,
            "SELECT" => consts::HW_SELECT,
            "RESET" => consts::HW_RESET,
            _ => panic!("Invalid hardware command: {}", hw),
        };
        self.write(value)
    }

    pub fn barcode(
        &mut self,
        code: &str,
        kind: &str,
        position: &str,
        font: &str,
        width: usize,
        height: usize,
    ) -> io::Result<usize> {
        let mut n = 0;
        if width >= 1 || width <= 255 {
            n += self.write(consts::BARCODE_WIDTH)?;
        }
        if height >= 2 || height <= 6 {
            n += self.write(consts::BARCODE_HEIGHT)?;
        }

        let font = font.to_uppercase();
        let position = position.to_uppercase();
        let kind = kind.to_uppercase().replace("-", "_");
        let font_value = match font.as_ref() {
            "B" => consts::BARCODE_FONT_B,
            "A" | _ => consts::BARCODE_FONT_A,
        };
        let txt_value = match position.as_ref() {
            "OFF" => consts::BARCODE_TXT_OFF,
            "ABV" => consts::BARCODE_TXT_ABV,
            "BTH" => consts::BARCODE_TXT_BTH,
            "BLW" | _ => consts::BARCODE_TXT_BLW,
        };
        let kind_value = match kind.as_ref() {
            "UPC_A" => consts::BARCODE_UPC_A,
            "UPC_E" => consts::BARCODE_UPC_E,
            "EAN8" => consts::BARCODE_EAN8,
            "CODE39" => consts::BARCODE_CODE39,
            "ITF" => consts::BARCODE_ITF,
            "NW7" => consts::BARCODE_NW7,
            "EAN13" | _ => consts::BARCODE_EAN13,
        };
        n += self.write(font_value)?;
        self.write(txt_value)?;
        self.write(kind_value)?;
        self.write(code.as_bytes())?;
        Ok(n)
    }

    #[cfg(feature = "qrcode")]
    pub fn qrimage(&mut self) -> io::Result<usize> {}

    #[cfg(feature = "qrcode")]
    pub fn qrcode(
        &mut self,
        code: &str,
        version: Option<i32>,
        level: &str,
        size: Option<i32>,
    ) -> io::Result<usize> {
        let level = level.to_uppercase();
        let level_value = match level.as_ref() {
            "M" => consts::QR_LEVEL_M,
            "Q" => consts::QR_LEVEL_Q,
            "H" => consts::QR_LEVEL_H,
            "L" | _ => consts::QR_LEVEL_L,
        };
        let mut n = 0;
        n += self.write(consts::TYPE_QR)?;
        n += self.write(consts::CODE2D)?;
        n += self.write_u8(version.unwrap_or(3) as u8)?;
        n += self.write(level_value)?;
        n += self.write_u8(size.unwrap_or(3) as u8)?;
        n += self.write_u16le(code.len() as u16)?;
        n += self.write(code.as_bytes())?;
        Ok(n)
    }

    pub fn cashdraw(&mut self, pin: i32) -> io::Result<usize> {
        let pin_value = if pin == 5 {
            consts::CD_KICK_5
        } else {
            consts::CD_KICK_2
        };
        self.write(pin_value)
    }

    pub fn cut(&mut self, part: bool) -> io::Result<usize> {
        let mut n_bytes = 0;
        n_bytes += self.print(
            iter::repeat(consts::EOL)
                .take(3)
                .collect::<String>()
                .as_ref(),
        )?;
        let paper_cut_type = if part {
            consts::PAPER_PART_CUT
        } else {
            consts::PAPER_FULL_CUT
        };
        n_bytes += self.write(paper_cut_type)?;
        Ok(n_bytes)
    }

    pub fn bit_image(&mut self, image: &Image, density: Option<&str>) -> io::Result<usize> {
        let density = density.unwrap_or("d24");
        let density_upper = density.to_uppercase();
        let header = match density_upper.as_ref() {
            "S8" => consts::BITMAP_S8,
            "D8" => consts::BITMAP_D8,
            "S24" => consts::BITMAP_S24,
            "D24" | _ => consts::BITMAP_D24,
        };
        let n = if density == "s8" || density == "d8" {
            1
        } else {
            3
        };
        let mut n_bytes = 0;
        n_bytes += self.line_space(0)?;
        for line in image.bitimage_lines(n * 8) {
            n_bytes += self.write(header)?;
            n_bytes += self.write_u16le((line.len() / n as usize) as u16)?;
            n_bytes += self.write(line.as_ref())?;
            n_bytes += self.feed(1)?;
        }
        Ok(n_bytes)
    }

    pub fn raster(&mut self, image: &Image, mode: Option<&str>) -> io::Result<usize> {
        let mode_upper = mode.unwrap_or("NORMAL").to_uppercase();
        let header = match mode_upper.as_ref() {
            "DH" => consts::GSV0_DH,
            "DWDH" => consts::GSV0_DWDH,
            "DW" => consts::GSV0_DW,
            "NORMAL" | _ => consts::GSV0_NORMAL,
        };
        let mut n_bytes = 0;
        n_bytes += self.write(header)?;
        n_bytes += self.write_u16le(((image.width+7)/8) as u16)?;
        n_bytes += self.write_u16le(image.height as u16)?;
        n_bytes += self.write(image.get_raster().as_ref())?;
        Ok(n_bytes)
    }
}
<commit_msg>Replaced printer.rs write panics with io::ErrorKind::InvalidData<commit_after>use std::io::{self, Write};
use std::iter;

use byteorder::{LittleEndian, WriteBytesExt};
use encoding::all::UTF_8;
use encoding::types::{EncoderTrap, EncodingRef};

use consts;
use img::Image;

pub struct Printer<W: io::Write> {
    writer: io::BufWriter<W>,
    codec: EncodingRef,
    trap: EncoderTrap,
}

impl<W: io::Write> Printer<W> {
    pub fn new(writer: W, codec: Option<EncodingRef>, trap: Option<EncoderTrap>) -> Printer<W> {
        Printer {
            writer: io::BufWriter::new(writer),
            codec: codec.unwrap_or(UTF_8 as EncodingRef),
            trap: trap.unwrap_or(EncoderTrap::Replace),
        }
    }

    fn encode(&mut self, content: &str) -> Vec<u8> {
        self.codec.encode(content, self.trap).unwrap()
    }

    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.writer.write(buf)
    }

    pub fn write_u8(&mut self, n: u8) -> io::Result<usize> {
        self.write(vec![n].as_slice())
    }

    fn write_u16le(&mut self, n: u16) -> io::Result<usize> {
        let mut wtr = vec![];
        wtr.write_u16::<LittleEndian>(n).unwrap();
        self.write(wtr.as_slice())
    }

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

    pub fn hwinit(&mut self) -> io::Result<usize> {
        self.write(consts::HW_INIT)
    }

    pub fn hwselect(&mut self) -> io::Result<usize> {
        self.write(consts::HW_SELECT)
    }

    pub fn hwreset(&mut self) -> io::Result<usize> {
        self.write(consts::HW_RESET)
    }

    pub fn print(&mut self, content: &str) -> io::Result<usize> {
        // let rv = self.encode(content);
        let rv = self.encode(content);
        self.write(rv.as_slice())
    }

    pub fn println(&mut self, content: &str) -> io::Result<usize> {
        self.print(format!("{}{}", content, consts::EOL).as_ref())
    }

    pub fn text(&mut self, content: &str) -> io::Result<usize> {
        self.println(content)
    }

    pub fn line_space(&mut self, n: i32) -> io::Result<usize> {
        if n >= 0 {
            Ok(self.write(consts::LS_SET)? + self.write_u8(n as u8)?)
        } else {
            self.write(consts::LS_DEFAULT)
        }
    }

    pub fn feed(&mut self, n: usize) -> io::Result<usize> {
        let n = if n < 1 { 1 } else { n };
        self.write(
            iter::repeat(consts::EOL)
                .take(n)
                .collect::<String>()
                .as_ref(),
        )
    }

    pub fn control(&mut self, ctrl: &str) -> io::Result<usize> {
        let ctrl_upper = ctrl.to_uppercase();
        let ctrl_value = match ctrl_upper.as_ref() {
            "LF" => consts::CTL_LF,
            "FF" => consts::CTL_FF,
            "CR" => consts::CTL_CR,
            "HT" => consts::CTL_HT,
            "VT" => consts::CTL_VT,
            _ => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("Invalid control action: {}", ctrl),
                ))
            }
        };
        self.write(ctrl_value)
    }

    pub fn align(&mut self, alignment: &str) -> io::Result<usize> {
        let align_upper = alignment.to_uppercase();
        let align_value = match align_upper.as_ref() {
            "LT" => consts::TXT_ALIGN_LT,
            "CT" => consts::TXT_ALIGN_CT,
            "RT" => consts::TXT_ALIGN_RT,
            _ => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("Invalid alignment: {}", alignment),
                ))
            }
        };
        self.write(align_value)
    }

    pub fn font(&mut self, family: &str) -> io::Result<usize> {
        let family_upper = family.to_uppercase();
        let family_value = match family_upper.as_ref() {
            "A" => consts::TXT_FONT_A,
            "B" => consts::TXT_FONT_B,
            "C" => consts::TXT_FONT_C,
            _ => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("Invalid font family: {}", family),
                ))
            }
        };
        self.write(family_value)
    }

    pub fn style(&mut self, kind: &str) -> io::Result<usize> {
        let kind_upper = kind.to_uppercase();
        match kind_upper.as_ref() {
            "B" => Ok(self.write(consts::TXT_UNDERL_OFF)? + self.write(consts::TXT_BOLD_ON)?),
            "U" => Ok(self.write(consts::TXT_BOLD_OFF)? + self.write(consts::TXT_UNDERL_ON)?),
            "U2" => Ok(self.write(consts::TXT_BOLD_OFF)? + self.write(consts::TXT_UNDERL2_ON)?),
            "BU" => Ok(self.write(consts::TXT_BOLD_ON)? + self.write(consts::TXT_UNDERL_ON)?),
            "BU2" => Ok(self.write(consts::TXT_BOLD_ON)? + self.write(consts::TXT_UNDERL2_ON)?),
            "NORMAL" | _ => {
                Ok(self.write(consts::TXT_BOLD_OFF)? + self.write(consts::TXT_UNDERL_OFF)?)
            }
        }
    }

    pub fn size(&mut self, width: usize, height: usize) -> io::Result<usize> {
        let mut n = self.write(consts::TXT_NORMAL)?;
        if width == 2 {
            n += self.write(consts::TXT_2WIDTH)?;
        }
        if height == 2 {
            n += self.write(consts::TXT_2HEIGHT)?;
        }
        Ok(n)
    }

    pub fn hardware(&mut self, hw: &str) -> io::Result<usize> {
        let value = match hw {
            "INIT" => consts::HW_INIT,
            "SELECT" => consts::HW_SELECT,
            "RESET" => consts::HW_RESET,
            _ => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("Invalid hardware command: {}", hw),
                ))
            }
        };
        self.write(value)
    }

    pub fn barcode(
        &mut self,
        code: &str,
        kind: &str,
        position: &str,
        font: &str,
        width: usize,
        height: usize,
    ) -> io::Result<usize> {
        let mut n = 0;
        if width >= 1 || width <= 255 {
            n += self.write(consts::BARCODE_WIDTH)?;
        }
        if height >= 2 || height <= 6 {
            n += self.write(consts::BARCODE_HEIGHT)?;
        }

        let font = font.to_uppercase();
        let position = position.to_uppercase();
        let kind = kind.to_uppercase().replace("-", "_");
        let font_value = match font.as_ref() {
            "B" => consts::BARCODE_FONT_B,
            "A" | _ => consts::BARCODE_FONT_A,
        };
        let txt_value = match position.as_ref() {
            "OFF" => consts::BARCODE_TXT_OFF,
            "ABV" => consts::BARCODE_TXT_ABV,
            "BTH" => consts::BARCODE_TXT_BTH,
            "BLW" | _ => consts::BARCODE_TXT_BLW,
        };
        let kind_value = match kind.as_ref() {
            "UPC_A" => consts::BARCODE_UPC_A,
            "UPC_E" => consts::BARCODE_UPC_E,
            "EAN8" => consts::BARCODE_EAN8,
            "CODE39" => consts::BARCODE_CODE39,
            "ITF" => consts::BARCODE_ITF,
            "NW7" => consts::BARCODE_NW7,
            "EAN13" | _ => consts::BARCODE_EAN13,
        };
        n += self.write(font_value)?;
        self.write(txt_value)?;
        self.write(kind_value)?;
        self.write(code.as_bytes())?;
        Ok(n)
    }

    #[cfg(feature = "qrcode")]
    pub fn qrimage(&mut self) -> io::Result<usize> {}

    #[cfg(feature = "qrcode")]
    pub fn qrcode(
        &mut self,
        code: &str,
        version: Option<i32>,
        level: &str,
        size: Option<i32>,
    ) -> io::Result<usize> {
        let level = level.to_uppercase();
        let level_value = match level.as_ref() {
            "M" => consts::QR_LEVEL_M,
            "Q" => consts::QR_LEVEL_Q,
            "H" => consts::QR_LEVEL_H,
            "L" | _ => consts::QR_LEVEL_L,
        };
        let mut n = 0;
        n += self.write(consts::TYPE_QR)?;
        n += self.write(consts::CODE2D)?;
        n += self.write_u8(version.unwrap_or(3) as u8)?;
        n += self.write(level_value)?;
        n += self.write_u8(size.unwrap_or(3) as u8)?;
        n += self.write_u16le(code.len() as u16)?;
        n += self.write(code.as_bytes())?;
        Ok(n)
    }

    pub fn cashdraw(&mut self, pin: i32) -> io::Result<usize> {
        let pin_value = if pin == 5 {
            consts::CD_KICK_5
        } else {
            consts::CD_KICK_2
        };
        self.write(pin_value)
    }

    pub fn cut(&mut self, part: bool) -> io::Result<usize> {
        let mut n_bytes = 0;
        n_bytes += self.print(
            iter::repeat(consts::EOL)
                .take(3)
                .collect::<String>()
                .as_ref(),
        )?;
        let paper_cut_type = if part {
            consts::PAPER_PART_CUT
        } else {
            consts::PAPER_FULL_CUT
        };
        n_bytes += self.write(paper_cut_type)?;
        Ok(n_bytes)
    }

    pub fn bit_image(&mut self, image: &Image, density: Option<&str>) -> io::Result<usize> {
        let density = density.unwrap_or("d24");
        let density_upper = density.to_uppercase();
        let header = match density_upper.as_ref() {
            "S8" => consts::BITMAP_S8,
            "D8" => consts::BITMAP_D8,
            "S24" => consts::BITMAP_S24,
            "D24" | _ => consts::BITMAP_D24,
        };
        let n = if density == "s8" || density == "d8" {
            1
        } else {
            3
        };
        let mut n_bytes = 0;
        n_bytes += self.line_space(0)?;
        for line in image.bitimage_lines(n * 8) {
            n_bytes += self.write(header)?;
            n_bytes += self.write_u16le((line.len() / n as usize) as u16)?;
            n_bytes += self.write(line.as_ref())?;
            n_bytes += self.feed(1)?;
        }
        Ok(n_bytes)
    }

    pub fn raster(&mut self, image: &Image, mode: Option<&str>) -> io::Result<usize> {
        let mode_upper = mode.unwrap_or("NORMAL").to_uppercase();
        let header = match mode_upper.as_ref() {
            "DH" => consts::GSV0_DH,
            "DWDH" => consts::GSV0_DWDH,
            "DW" => consts::GSV0_DW,
            "NORMAL" | _ => consts::GSV0_NORMAL,
        };
        let mut n_bytes = 0;
        n_bytes += self.write(header)?;
        n_bytes += self.write_u16le(((image.width + 7) / 8) as u16)?;
        n_bytes += self.write_u16le(image.height as u16)?;
        n_bytes += self.write(image.get_raster().as_ref())?;
        Ok(n_bytes)
    }
}
<|endoftext|>


<commit_before>use std;
use variables::*;
use variables::LpExpression::*;
use std::rc::Rc;
use std::collections::HashMap;
use std::fs::File;
use std::io::prelude::*;

//use variables::LpExpression::{AddExpr, MulExpr};
use std::ops::{AddAssign};

//use std::collections::HashMap;

/// Enum helping to specify the objective function of the linear problem.
///
/// # Examples:
///
/// ```
/// use lp_modeler::problem::{LpObjective, LpProblem};
///
/// let mut problem = LpProblem::new("One Problem", LpObjective::Maximize);
/// ```
#[derive(Debug, PartialEq)]
pub enum LpObjective {
    Minimize,
    Maximize
}

pub trait Problem {
    fn add_objective_expression(&mut self, expr: &LpExpression);
    fn add_constraints(&mut self, expr: &LpConstraint);
}

/// Structure used for creating the model and solving a linear problem.
///
/// # Examples:
///
/// ```
/// use lp_modeler::problem::{LpObjective, Problem, LpProblem};
/// use lp_modeler::operations::{LpOperations};
/// use lp_modeler::variables::LpInteger;
/// use lp_modeler::solvers::{SolverTrait, CbcSolver};
///
/// let ref a = LpInteger::new("a");
/// let ref b = LpInteger::new("b");
/// let ref c = LpInteger::new("c");
///
/// let mut problem = LpProblem::new("One Problem", LpObjective::Maximize);
/// problem += 10.0 * a + 20.0 * b;
///
/// problem += (500*a + 1200*b + 1500*c).le(10000);
/// problem += (a + b*2 + c).le(10);
/// problem += (a).le(b);
///
/// let solver = CbcSolver::new();
///
/// match solver.run(&problem) {
/// Ok((status, res)) => {
///     println!("Status {:?}", status);
///         for (name, value) in res.iter() {
///             println!("value of {} = {}", name, value);
///         }
///     },
///     Err(msg) => println!("{}", msg),
/// }
/// ```
#[derive(Debug)]
pub struct LpProblem {
    name: &'static str,
    objective_type: LpObjective,
    obj_expr: Option<LpExpression>,
    constraints: Vec<LpConstraint>

}

impl LpProblem {
    /// Create a new problem
    pub fn new(name: &'static str, objective: LpObjective) -> LpProblem {
        LpProblem { name: name, objective_type: objective, obj_expr: None, constraints: Vec::new() }
    }

    // TODO: Call once and pass into parameter
    // TODO: Check variables on the objective function
    fn variables(&self) -> HashMap<String, &LpExpression> {
        fn var<'a>(expr: &'a LpExpression, lst: &mut Vec<(String, &'a LpExpression)>) {
            match expr {
                &ConsBin(LpBinary { ref name, .. }) |
                &ConsInt(LpInteger { ref name, .. }) |
                &ConsCont(LpContinuous { ref name, .. })
                => { lst.push((name.clone(), expr)); },

                &MulExpr(_, ref e) => { var(&*e, lst); },
                &AddExpr(ref e1, ref e2) | &SubExpr(ref e1, ref e2) => {
                    var(&*e1, lst);
                    var(&*e2, lst);
                },
                _ => ()
            }
        }

        let mut lst: Vec<_> = Vec::new();
        for e in &self.constraints {
            var(&e.0, &mut lst);
        }
        lst.iter().map(|&(ref n, ref x)| (n.clone(), *x)).collect::<HashMap<String, &LpExpression>>()
    }

}

pub trait LpFileFormat {
    fn to_lp_file_format(&self) -> String;
    fn write_lp(&self, file_model: &str) -> std::io::Result<()>  {
        let mut buffer = try!(File::create(file_model));
        try!(buffer.write(self.to_lp_file_format().as_bytes()));
        Ok(())
    }
}

impl LpFileFormat for LpProblem {
    fn to_lp_file_format(&self) -> String {

        let objective_string = || {
            if let Some(ref expr) = self.obj_expr {
                expr.to_lp_file_format()
            } else {
                String::new()
            }
        };

        let constraints_string = || {
            let mut res = String::new();
            let mut cstrs = self.constraints.iter();
            let mut index = 1;
            while let Some(ref constraint) = cstrs.next() {
                res.push_str("  c");
                res.push_str(&index.to_string());
                res.push_str(": ");
                index += 1;

                res.push_str(&constraint.to_lp_file_format());

                res.push_str("\n");
            }
            res
        };

        let bounds_string = || {
            let mut res = String::new();
            for (_, v) in self.variables() {
                match v {
                    &ConsInt(LpInteger { ref name, lower_bound, upper_bound })
                    | &ConsCont(LpContinuous { ref name, lower_bound, upper_bound }) => {
                        if let Some(l) = lower_bound {
                            res.push_str("  ");
                            res.push_str(&l.to_string());
                            res.push_str(" <= ");
                            res.push_str(&name);
                            if let Some(u) = upper_bound {
                                res.push_str(" <= ");
                                res.push_str(&u.to_string());
                            }
                            res.push_str("\n");
                        } else if let Some(u) = upper_bound {
                            res.push_str("  ");
                            res.push_str(&name);
                            res.push_str(" <= ");
                            res.push_str(&u.to_string());
                            res.push_str("\n");
                        } else {
                            match v {
                                &ConsCont(LpContinuous { .. }) => {
                                    res.push_str("  ");
                                    res.push_str(&name);
                                    res.push_str(" free\n");
                                }, // TODO: IntegerVar => -INF to INF
                                _ => ()
                            }
                        }
                    },
                    _ => (),
                }
            }
            res
        };

        let integers_string = || {
            let mut res = String::new();
            for (_, v) in self.variables() {
                match v {
                    &ConsInt(LpInteger { ref name, .. }) => {
                        res.push_str(name);
                        res.push_str(" ");
                    },
                    _ => (),
                }
            }
            res
        };

        let binaries_string = || {
            let mut res = String::new();
            for (_, v) in self.variables() {
                match v {
                    &ConsBin(LpBinary { ref name }) => {
                        res.push_str(name);
                        res.push_str(" ");
                    },
                    _ => (),
                }
            }
            res
        };

        let mut buffer = String::new();

        buffer.push_str("\\ ");
        buffer.push_str(&self.name);
        buffer.push_str("\n\n");

        // Write objectives
        match self.objective_type {
            LpObjective::Maximize => { buffer.push_str("Maximize\n  "); },
            LpObjective::Minimize => { buffer.push_str("Minimize\n  "); },
        }
        let obj_str = objective_string();
        buffer.push_str(&obj_str);

        // Write constraints
        let cstr_str = constraints_string();
        if cstr_str.len() > 0 {
            buffer.push_str("\n\nSubject To\n");
            buffer.push_str(&cstr_str);
        }

        // Write bounds for Integer and Continuous
        let bound_str = bounds_string();
        if bound_str.len() > 0 {
            buffer.push_str("\nBounds\n");
            buffer.push_str(&bound_str);
        }

        // Write Integer vars
        let generals_str = integers_string();
        if generals_str.len() > 0 {
            buffer.push_str("\nGenerals\n  ");
            buffer.push_str(&generals_str);
            buffer.push_str("\n");
        }

        // Write Binaries vars
        let binaries_str = binaries_string();
        if binaries_str.len() > 0 {
            buffer.push_str("\nBinary\n  ");
            buffer.push_str(&binaries_str);
            buffer.push_str("\n");
        }

        buffer.push_str("\nEnd");

        buffer
    }
}

impl Problem for LpProblem {

    fn add_objective_expression(&mut self, expr: &LpExpression) {
        if let Some(e) = self.obj_expr.clone() {
            self.obj_expr = Some(AddExpr(Rc::new(expr.clone()), Rc::new(e.clone())).dfs_remove_constant());
        } else {
            self.obj_expr = Some(expr.dfs_remove_constant());
        }
    }

    fn add_constraints(&mut self, expr: &LpConstraint) {
        self.constraints.push(expr.clone());
    }

}

/*
/// Add constraints
impl AddAssign<LpConstraint> for LpProblem {
    fn add_assign(&mut self, _rhs: LpConstraint) {
        self.add_constraints(&_rhs);
    }
}

/// Add an expression as an objective function
impl<T> AddAssign<T> for LpProblem where T: Into<LpExpression>{
    fn add_assign(&mut self, _rhs: T) {
        self.add_objective_expression(&_rhs.into());
    }
}
*/

macro_rules! impl_addassign_for_generic_problem {
    ($problem: ty) => {
        /// Add constraints
        impl AddAssign<LpConstraint> for $problem {
            fn add_assign(&mut self, _rhs: LpConstraint) {
                self.add_constraints(&_rhs);
            }
        }
        /// Add an expression as an objective function
        impl<T> AddAssign<T> for $problem where T: Into<LpExpression>{
            fn add_assign(&mut self, _rhs: T) {
                self.add_objective_expression(&_rhs.into());
            }
        }
    }
}
impl_addassign_for_generic_problem!(LpProblem);
<commit_msg>Add a name to the objective function in LP files<commit_after>use std;
use variables::*;
use variables::LpExpression::*;
use std::rc::Rc;
use std::collections::HashMap;
use std::fs::File;
use std::io::prelude::*;

//use variables::LpExpression::{AddExpr, MulExpr};
use std::ops::{AddAssign};

//use std::collections::HashMap;

/// Enum helping to specify the objective function of the linear problem.
///
/// # Examples:
///
/// ```
/// use lp_modeler::problem::{LpObjective, LpProblem};
///
/// let mut problem = LpProblem::new("One Problem", LpObjective::Maximize);
/// ```
#[derive(Debug, PartialEq)]
pub enum LpObjective {
    Minimize,
    Maximize
}

pub trait Problem {
    fn add_objective_expression(&mut self, expr: &LpExpression);
    fn add_constraints(&mut self, expr: &LpConstraint);
}

/// Structure used for creating the model and solving a linear problem.
///
/// # Examples:
///
/// ```
/// use lp_modeler::problem::{LpObjective, Problem, LpProblem};
/// use lp_modeler::operations::{LpOperations};
/// use lp_modeler::variables::LpInteger;
/// use lp_modeler::solvers::{SolverTrait, CbcSolver};
///
/// let ref a = LpInteger::new("a");
/// let ref b = LpInteger::new("b");
/// let ref c = LpInteger::new("c");
///
/// let mut problem = LpProblem::new("One Problem", LpObjective::Maximize);
/// problem += 10.0 * a + 20.0 * b;
///
/// problem += (500*a + 1200*b + 1500*c).le(10000);
/// problem += (a + b*2 + c).le(10);
/// problem += (a).le(b);
///
/// let solver = CbcSolver::new();
///
/// match solver.run(&problem) {
/// Ok((status, res)) => {
///     println!("Status {:?}", status);
///         for (name, value) in res.iter() {
///             println!("value of {} = {}", name, value);
///         }
///     },
///     Err(msg) => println!("{}", msg),
/// }
/// ```
#[derive(Debug)]
pub struct LpProblem {
    name: &'static str,
    objective_type: LpObjective,
    obj_expr: Option<LpExpression>,
    constraints: Vec<LpConstraint>

}

impl LpProblem {
    /// Create a new problem
    pub fn new(name: &'static str, objective: LpObjective) -> LpProblem {
        LpProblem { name: name, objective_type: objective, obj_expr: None, constraints: Vec::new() }
    }

    // TODO: Call once and pass into parameter
    // TODO: Check variables on the objective function
    fn variables(&self) -> HashMap<String, &LpExpression> {
        fn var<'a>(expr: &'a LpExpression, lst: &mut Vec<(String, &'a LpExpression)>) {
            match expr {
                &ConsBin(LpBinary { ref name, .. }) |
                &ConsInt(LpInteger { ref name, .. }) |
                &ConsCont(LpContinuous { ref name, .. })
                => { lst.push((name.clone(), expr)); },

                &MulExpr(_, ref e) => { var(&*e, lst); },
                &AddExpr(ref e1, ref e2) | &SubExpr(ref e1, ref e2) => {
                    var(&*e1, lst);
                    var(&*e2, lst);
                },
                _ => ()
            }
        }

        let mut lst: Vec<_> = Vec::new();
        for e in &self.constraints {
            var(&e.0, &mut lst);
        }
        lst.iter().map(|&(ref n, ref x)| (n.clone(), *x)).collect::<HashMap<String, &LpExpression>>()
    }

}

pub trait LpFileFormat {
    fn to_lp_file_format(&self) -> String;
    fn write_lp(&self, file_model: &str) -> std::io::Result<()>  {
        let mut buffer = try!(File::create(file_model));
        try!(buffer.write(self.to_lp_file_format().as_bytes()));
        Ok(())
    }
}

impl LpFileFormat for LpProblem {
    fn to_lp_file_format(&self) -> String {

        let objective_string = || {
            if let Some(ref expr) = self.obj_expr {
                format!("obj: {}", expr.to_lp_file_format())
            } else {
                String::new()
            }
        };

        let constraints_string = || {
            let mut res = String::new();
            let mut cstrs = self.constraints.iter();
            let mut index = 1;
            while let Some(ref constraint) = cstrs.next() {
                res.push_str("  c");
                res.push_str(&index.to_string());
                res.push_str(": ");
                index += 1;

                res.push_str(&constraint.to_lp_file_format());

                res.push_str("\n");
            }
            res
        };

        let bounds_string = || {
            let mut res = String::new();
            for (_, v) in self.variables() {
                match v {
                    &ConsInt(LpInteger { ref name, lower_bound, upper_bound })
                    | &ConsCont(LpContinuous { ref name, lower_bound, upper_bound }) => {
                        if let Some(l) = lower_bound {
                            res.push_str("  ");
                            res.push_str(&l.to_string());
                            res.push_str(" <= ");
                            res.push_str(&name);
                            if let Some(u) = upper_bound {
                                res.push_str(" <= ");
                                res.push_str(&u.to_string());
                            }
                            res.push_str("\n");
                        } else if let Some(u) = upper_bound {
                            res.push_str("  ");
                            res.push_str(&name);
                            res.push_str(" <= ");
                            res.push_str(&u.to_string());
                            res.push_str("\n");
                        } else {
                            match v {
                                &ConsCont(LpContinuous { .. }) => {
                                    res.push_str("  ");
                                    res.push_str(&name);
                                    res.push_str(" free\n");
                                }, // TODO: IntegerVar => -INF to INF
                                _ => ()
                            }
                        }
                    },
                    _ => (),
                }
            }
            res
        };

        let integers_string = || {
            let mut res = String::new();
            for (_, v) in self.variables() {
                match v {
                    &ConsInt(LpInteger { ref name, .. }) => {
                        res.push_str(name);
                        res.push_str(" ");
                    },
                    _ => (),
                }
            }
            res
        };

        let binaries_string = || {
            let mut res = String::new();
            for (_, v) in self.variables() {
                match v {
                    &ConsBin(LpBinary { ref name }) => {
                        res.push_str(name);
                        res.push_str(" ");
                    },
                    _ => (),
                }
            }
            res
        };

        let mut buffer = String::new();

        buffer.push_str("\\ ");
        buffer.push_str(&self.name);
        buffer.push_str("\n\n");

        // Write objectives
        match self.objective_type {
            LpObjective::Maximize => { buffer.push_str("Maximize\n  "); },
            LpObjective::Minimize => { buffer.push_str("Minimize\n  "); },
        }
        let obj_str = objective_string();
        buffer.push_str(&obj_str);

        // Write constraints
        let cstr_str = constraints_string();
        if cstr_str.len() > 0 {
            buffer.push_str("\n\nSubject To\n");
            buffer.push_str(&cstr_str);
        }

        // Write bounds for Integer and Continuous
        let bound_str = bounds_string();
        if bound_str.len() > 0 {
            buffer.push_str("\nBounds\n");
            buffer.push_str(&bound_str);
        }

        // Write Integer vars
        let generals_str = integers_string();
        if generals_str.len() > 0 {
            buffer.push_str("\nGenerals\n  ");
            buffer.push_str(&generals_str);
            buffer.push_str("\n");
        }

        // Write Binaries vars
        let binaries_str = binaries_string();
        if binaries_str.len() > 0 {
            buffer.push_str("\nBinary\n  ");
            buffer.push_str(&binaries_str);
            buffer.push_str("\n");
        }

        buffer.push_str("\nEnd");

        buffer
    }
}

impl Problem for LpProblem {

    fn add_objective_expression(&mut self, expr: &LpExpression) {
        if let Some(e) = self.obj_expr.clone() {
            self.obj_expr = Some(AddExpr(Rc::new(expr.clone()), Rc::new(e.clone())).dfs_remove_constant());
        } else {
            self.obj_expr = Some(expr.dfs_remove_constant());
        }
    }

    fn add_constraints(&mut self, expr: &LpConstraint) {
        self.constraints.push(expr.clone());
    }

}

/*
/// Add constraints
impl AddAssign<LpConstraint> for LpProblem {
    fn add_assign(&mut self, _rhs: LpConstraint) {
        self.add_constraints(&_rhs);
    }
}

/// Add an expression as an objective function
impl<T> AddAssign<T> for LpProblem where T: Into<LpExpression>{
    fn add_assign(&mut self, _rhs: T) {
        self.add_objective_expression(&_rhs.into());
    }
}
*/

macro_rules! impl_addassign_for_generic_problem {
    ($problem: ty) => {
        /// Add constraints
        impl AddAssign<LpConstraint> for $problem {
            fn add_assign(&mut self, _rhs: LpConstraint) {
                self.add_constraints(&_rhs);
            }
        }
        /// Add an expression as an objective function
        impl<T> AddAssign<T> for $problem where T: Into<LpExpression>{
            fn add_assign(&mut self, _rhs: T) {
                self.add_objective_expression(&_rhs.into());
            }
        }
    }
}
impl_addassign_for_generic_problem!(LpProblem);
<|endoftext|>


<commit_before>use threadpool::ThreadPool;

use std::process::{Child, Command};
use std::sync::mpsc::{Sender};

pub struct Process {
    process: Child,
    killed: bool
}

#[cfg(target_family = "unix")]
impl Process {
    pub fn new(cmd: &str, updated_paths: Vec<&str>) -> Option<Process>{
        use libc;
        use std::os::unix::process::CommandExt;

        let mut command = Command::new("sh");
        command.arg("-c").arg(cmd);

        if !updated_paths.is_empty() {
            command.env("WATCHEXEC_UPDATED_PATH", updated_paths[0]);
        }

        command.before_exec(|| unsafe {
            libc::setpgid(0, 0);
            Ok(())
        })
        .spawn()
        .ok()
        .and_then(|p| Some(Process { process: p, killed: false }))
    }

    pub fn kill(&mut self) {
        if self.killed {
            return;
        }

        use libc;

        extern "C" {
            fn killpg(pgrp: libc::pid_t, sig: libc::c_int) -> libc::c_int;
        }

        unsafe {
            killpg(self.process.id() as i32, libc::SIGTERM);
        }

        self.killed = true;
    }

    pub fn wait(&mut self) {
        use nix::sys::wait::waitpid;

        let pid = self.process.id() as i32;
        let _ = waitpid(-pid, None);
    }
}

#[cfg(target_family = "windows")]
impl Process {
    pub fn new(cmd: &str, updated_paths: Vec<&str>) -> Option<Process> {
        use std::os::windows::io::AsRawHandle;

        let mut command = Command::new("cmd.exe");
        command.arg("/C").arg(cmd);

        if !updated_paths.is_empty() {
            command.env("WATCHEXEC_UPDATED_PATH", updated_paths[0]);
        }

        command.spawn()
            .ok()
            .and_then(|p| { Some(Process { process: p, killed: false })})
    }

    pub fn kill(&mut self) {
        if self.killed {
            return;
        }

        self.process.kill();
        self.killed = true;
    }

    pub fn wait(&mut self) {
        self.process.wait();
    }
}

impl Drop for Process {
    fn drop(&mut self) {
        self.kill();
    }
}

pub struct ProcessReaper {
    pool: ThreadPool,
    tx: Sender<()>,
}

impl ProcessReaper {
    pub fn new(tx: Sender<()>) -> ProcessReaper {
        ProcessReaper {
            pool: ThreadPool::new(1),
            tx: tx
        }
    }

    pub fn wait_process(&self, mut process: Process) {
        let tx = self.tx.clone();

        self.pool.execute(move || {
            process.wait();
            let _ = tx.send(());
        });
    }
}

#[cfg(test)]
#[cfg(target_family = "unix")]
mod tests {
    use std::path::Path;
    use std::thread;
    use std::time::Duration;

    use mktemp::Temp;

    use super::Process;

    fn file_contents(path: &Path) -> String {
        use std::fs::File;
        use std::io::Read;

        let mut f = File::open(path).unwrap();
        let mut s = String::new();
        f.read_to_string(&mut s).unwrap();

        s
    }

    #[test]
    fn test_start() {
        let process = Process::new("echo hi", vec![]);

        assert!(process.is_some());
    }

    #[test]
    fn test_wait() {
        let file = Temp::new_file().unwrap();
        let path = file.to_path_buf();
        let mut process = Process::new(&format!("echo hi > {}", path.to_str().unwrap()), vec![]).unwrap();
        process.wait();

        assert!(file_contents(&path).starts_with("hi"));
    }

    #[test]
    fn test_kill() {
        let file = Temp::new_file().unwrap();
        let path = file.to_path_buf();

        let mut process = Process::new(&format!("sleep 20; echo hi > {}", path.to_str().unwrap()), vec![]).unwrap();
        thread::sleep(Duration::from_millis(250));
        process.kill();
        process.wait();

        assert!(file_contents(&path) == "");
    }
}
<commit_msg>Fix Windows warnings<commit_after>use threadpool::ThreadPool;

use std::process::{Child, Command};
use std::sync::mpsc::{Sender};

pub struct Process {
    process: Child,
    killed: bool
}

#[cfg(target_family = "unix")]
impl Process {
    pub fn new(cmd: &str, updated_paths: Vec<&str>) -> Option<Process>{
        use libc;
        use std::os::unix::process::CommandExt;

        let mut command = Command::new("sh");
        command.arg("-c").arg(cmd);

        if !updated_paths.is_empty() {
            command.env("WATCHEXEC_UPDATED_PATH", updated_paths[0]);
        }

        command.before_exec(|| unsafe {
            libc::setpgid(0, 0);
            Ok(())
        })
        .spawn()
        .ok()
        .and_then(|p| Some(Process { process: p, killed: false }))
    }

    pub fn kill(&mut self) {
        if self.killed {
            return;
        }

        use libc;

        extern "C" {
            fn killpg(pgrp: libc::pid_t, sig: libc::c_int) -> libc::c_int;
        }

        unsafe {
            killpg(self.process.id() as i32, libc::SIGTERM);
        }

        self.killed = true;
    }

    pub fn wait(&mut self) {
        use nix::sys::wait::waitpid;

        let pid = self.process.id() as i32;
        let _ = waitpid(-pid, None);
    }
}

#[cfg(target_family = "windows")]
impl Process {
    pub fn new(cmd: &str, updated_paths: Vec<&str>) -> Option<Process> {
        let mut command = Command::new("cmd.exe");
        command.arg("/C").arg(cmd);

        if !updated_paths.is_empty() {
            command.env("WATCHEXEC_UPDATED_PATH", updated_paths[0]);
        }

        command.spawn()
            .ok()
            .and_then(|p| { Some(Process { process: p, killed: false })})
    }

    pub fn kill(&mut self) {
        if self.killed {
            return;
        }

        let _ = self.process.kill();
        self.killed = true;
    }

    pub fn wait(&mut self) {
        let _ = self.process.wait();
    }
}

impl Drop for Process {
    fn drop(&mut self) {
        self.kill();
    }
}

pub struct ProcessReaper {
    pool: ThreadPool,
    tx: Sender<()>,
}

impl ProcessReaper {
    pub fn new(tx: Sender<()>) -> ProcessReaper {
        ProcessReaper {
            pool: ThreadPool::new(1),
            tx: tx
        }
    }

    pub fn wait_process(&self, mut process: Process) {
        let tx = self.tx.clone();

        self.pool.execute(move || {
            process.wait();
            let _ = tx.send(());
        });
    }
}

#[cfg(test)]
#[cfg(target_family = "unix")]
mod tests {
    use std::path::Path;
    use std::thread;
    use std::time::Duration;

    use mktemp::Temp;

    use super::Process;

    fn file_contents(path: &Path) -> String {
        use std::fs::File;
        use std::io::Read;

        let mut f = File::open(path).unwrap();
        let mut s = String::new();
        f.read_to_string(&mut s).unwrap();

        s
    }

    #[test]
    fn test_start() {
        let process = Process::new("echo hi", vec![]);

        assert!(process.is_some());
    }

    #[test]
    fn test_wait() {
        let file = Temp::new_file().unwrap();
        let path = file.to_path_buf();
        let mut process = Process::new(&format!("echo hi > {}", path.to_str().unwrap()), vec![]).unwrap();
        process.wait();

        assert!(file_contents(&path).starts_with("hi"));
    }

    #[test]
    fn test_kill() {
        let file = Temp::new_file().unwrap();
        let path = file.to_path_buf();

        let mut process = Process::new(&format!("sleep 20; echo hi > {}", path.to_str().unwrap()), vec![]).unwrap();
        thread::sleep(Duration::from_millis(250));
        process.kill();
        process.wait();

        assert!(file_contents(&path) == "");
    }
}
<|endoftext|>


<commit_before>extern crate rustc_serialize;

use std::cmp::min;
use std::cmp::Ordering;
use std::collections::HashMap;
use std::hash::Hash;
use std::io::prelude::*;
use std::fs::File;

use self::rustc_serialize::json;

use ngrams::ngrams;

pub type StringCountPair = (String, usize);
pub struct Profile {
    pub ngram_ranks: HashMap<String, usize>,
}

impl Profile {
    pub fn build_from_text(text: &str) -> Profile {
        let mut ngram_counts = HashMap::new();
        for n in 1..6 {
            for ngram in ngrams(text, n) {
                *ngram_counts.entry(ngram).or_insert(0usize) += 1;
            }
        }

        let mut ngrams_and_counts = vec_from_hashmap(&ngram_counts);
        ngrams_and_counts.sort_by(cmp_counts_reverse);

        let mut ngram_ranks = HashMap::new();
        let size = min(ngrams_and_counts.len(), 300);
        for (index, item) in ngrams_and_counts[..size].iter().enumerate() {
            ngram_ranks.insert((*item).clone().0, index);
        }

        Profile { ngram_ranks: ngram_ranks }
    }

    pub fn load_from_file(path: &str) -> Profile {
        let mut f = File::open(path).unwrap();
        let mut encoded_profile = String::new();
        f.read_to_string(&mut encoded_profile);
        let ngram_ranks = json::decode(&encoded_profile).unwrap();
        Profile { ngram_ranks: ngram_ranks }
    }

    pub fn save_to_file(&self, path: &str) {
        let encoded_profile = json::encode(&self.ngram_ranks).unwrap();
        let mut f = File::create(path).unwrap();
        f.write_all(encoded_profile.as_bytes()).unwrap();
    }

    pub fn compare(&self, other: &Profile) -> usize {
        let max_difference = other.ngram_ranks.len();
        let mut difference = 0;
        for (ngram, index) in self.ngram_ranks.iter() {
            difference += match other.ngram_ranks.get(ngram) {
                Some(other_index) => get_difference(*index, *other_index),
                None => max_difference,
            }
        }
        difference
    }
}

// useful utils

fn get_difference(a: usize, b: usize) -> usize {
    if a > b { a - b } else { b - a }
}

fn vec_from_hashmap<K, V> (hashmap: &HashMap<K, V>) -> Vec<(K, V)>
         where K: Eq + Hash + Clone, V: Clone {
    let mut vec: Vec<(K, V)> = Vec::new();
    for (key, value) in hashmap.iter() {
        let item = ((*key).clone(), (*value).clone());
        vec.push(item);
    }
    vec
}

fn cmp_counts_reverse(a: &StringCountPair, b: &StringCountPair) -> Ordering {
    b.1.cmp(&a.1)
}
<commit_msg>Add choosing best match among profiles<commit_after>extern crate rustc_serialize;

use std::cmp::min;
use std::cmp::Ordering;
use std::collections::HashMap;
use std::hash::Hash;
use std::io::prelude::*;
use std::fs::File;

use self::rustc_serialize::json;

use ngrams::ngrams;

pub type StringCountPair = (String, usize);
pub struct Profile {
    pub ngram_ranks: HashMap<String, usize>,
}

impl Profile {
    pub fn build_from_text(text: &str) -> Profile {
        let mut ngram_counts = HashMap::new();
        for n in 1..6 {
            for ngram in ngrams(text, n) {
                *ngram_counts.entry(ngram).or_insert(0usize) += 1;
            }
        }

        let mut ngrams_and_counts = vec_from_hashmap(&ngram_counts);
        ngrams_and_counts.sort_by(cmp_counts_reverse);

        let mut ngram_ranks = HashMap::new();
        let size = min(ngrams_and_counts.len(), 300);
        for (index, item) in ngrams_and_counts[..size].iter().enumerate() {
            ngram_ranks.insert((*item).clone().0, index);
        }

        Profile { ngram_ranks: ngram_ranks }
    }

    pub fn load_from_file(path: &str) -> Profile {
        let mut f = File::open(path).unwrap();
        let mut encoded_profile = String::new();
        f.read_to_string(&mut encoded_profile);
        let ngram_ranks = json::decode(&encoded_profile).unwrap();
        Profile { ngram_ranks: ngram_ranks }
    }

    pub fn save_to_file(&self, path: &str) {
        let encoded_profile = json::encode(&self.ngram_ranks).unwrap();
        let mut f = File::create(path).unwrap();
        f.write_all(encoded_profile.as_bytes()).unwrap();
    }

    pub fn choose_best_match<'a>(&'a self, others: &Vec<&'a Profile>) -> &Profile {
        find_min_by_key(others, |other| self.compare(other)).unwrap()
    }

    pub fn compare(&self, other: &Profile) -> usize {
        let max_difference = other.ngram_ranks.len();
        let mut difference = 0;
        for (ngram, index) in self.ngram_ranks.iter() {
            difference += match other.ngram_ranks.get(ngram) {
                Some(other_index) => get_difference(*index, *other_index),
                None => max_difference,
            }
        }
        difference
    }
}

// useful utils

fn find_min_by_key<T, I, F>(iterable: I, key: F) -> Option<T>
    where I : IntoIterator<Item=T>,
          F : Fn(&T) -> usize {

    let mut min_value = None;
    let mut min_item = None;
    for item in iterable {
        let value = key(&item);
        if min_value == None || value < min_value.unwrap() {
            min_value = Some(value);
            min_item = Some(item);
        }
    }
    min_item
}

fn get_difference(a: usize, b: usize) -> usize {
    if a > b { a - b } else { b - a }
}

fn vec_from_hashmap<K, V> (hashmap: &HashMap<K, V>) -> Vec<(K, V)>
         where K: Eq + Hash + Clone, V: Clone {
    let mut vec: Vec<(K, V)> = Vec::new();
    for (key, value) in hashmap.iter() {
        let item = ((*key).clone(), (*value).clone());
        vec.push(item);
    }
    vec
}

fn cmp_counts_reverse(a: &StringCountPair, b: &StringCountPair) -> Ordering {
    b.1.cmp(&a.1)
}
<|endoftext|>


<commit_before>use super::super::services::collection_service;
use super::super::rocket_contrib::JSON;
use super::super::models::{Collection, NewCollection};
use std::error::Error;
use rocket::response::status;

#[get("/collections")]
fn get_collections() -> Result<JSON<Vec<Collection>>, Box<Error>> {
    let result = try!(collection_service::get_all());
    Ok(JSON(result))
}

#[get("/collections/<id>")]
fn get_collection(id: i32) -> Option<JSON<Collection>> {
    match collection_service::find(id) {
        Ok(t) => Some(JSON(t)),
        Err(_) => None,
    }
}

#[post("/collections/new?<new_collection>")]
fn new(new_collection: NewCollection) -> Result<status::Created<JSON<Collection>>, Box<Error>> {
    let result = try!(collection_service::create(&new_collection));
    Ok(status::Created(
        format!("/collections/{}", result.id),
        Some(JSON(result)),
    ))
}

#[put("/collections/update?<collection>")]
fn update(collection: Collection) -> Option<JSON<Collection>> {
    match collection_service::update(&collection) {
        Ok(t) => Some(JSON(t)),
        Err(_) => None,
    }
}

#[post("/images/<image_id>/collections/<collection_id>")]
fn new_image_collection(image_id: i32, collection_id: i32) -> Option<status::NoContent> {
    match collection_service::add_image_to_collection(image_id, collection_id) {
        Ok(result) => Some(status::NoContent),
        Err(_) => None,
    }
}
<commit_msg>Change route def<commit_after>use super::super::services::collection_service;
use super::super::rocket_contrib::JSON;
use super::super::models::{Collection, NewCollection};
use std::error::Error;
use rocket::response::status;

#[get("/collections")]
fn get_collections() -> Result<JSON<Vec<Collection>>, Box<Error>> {
    let result = try!(collection_service::get_all());
    Ok(JSON(result))
}

#[get("/collections/<id>")]
fn get_collection(id: i32) -> Option<JSON<Collection>> {
    match collection_service::find(id) {
        Ok(t) => Some(JSON(t)),
        Err(_) => None,
    }
}

#[post("/collections/new?<new_collection>")]
fn new(new_collection: NewCollection) -> Result<status::Created<JSON<Collection>>, Box<Error>> {
    let result = try!(collection_service::create(&new_collection));
    Ok(status::Created(
        format!("/collections/{}", result.id),
        Some(JSON(result)),
    ))
}

#[put("/collections/update?<collection>")]
fn update(collection: Collection) -> Option<JSON<Collection>> {
    match collection_service::update(&collection) {
        Ok(t) => Some(JSON(t)),
        Err(_) => None,
    }
}

#[post("/collections/<collection_id>/images/<image_id>")]
fn new_image_collection(image_id: i32, collection_id: i32) -> Option<status::NoContent> {
    match collection_service::add_image_to_collection(image_id, collection_id) {
        Ok(result) => Some(status::NoContent),
        Err(_) => None,
    }
}
<|endoftext|>