<commit_before>use byteorder::{LittleEndian, ReadBytesExt};
use std::convert::TryFrom;
use std::io::{self, Cursor, Read, Seek, SeekFrom};
use std::marker::PhantomData;
use std::{error, fmt, mem};

use crate::color::ColorType;
use crate::error::{DecodingError, ImageError, ImageResult, UnsupportedError, UnsupportedErrorKind};
use crate::image::{self, ImageDecoder, ImageFormat};

use self::InnerDecoder::*;
use crate::bmp::BmpDecoder;
use crate::png::PngDecoder;

// http://www.w3.org/TR/PNG-Structure.html
// The first eight bytes of a PNG file always contain the following (decimal) values:
const PNG_SIGNATURE: [u8; 8] = [137, 80, 78, 71, 13, 10, 26, 10];

/// Errors that can occur during decoding and parsing an ICO image or one of its enclosed images.
#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
enum DecoderError {
    /// The ICO directory is empty
    NoEntries,
    /// The number of color planes (0 or 1), or the horizontal coordinate of the hotspot for CUR files too big.
    IcoEntryTooManyPlanesOrHotspot,
    /// The bit depth (may be 0 meaning unspecified), or the vertical coordinate of the hotspot for CUR files too big.
    IcoEntryTooManyBitsPerPixelOrHotspot,

    /// The entry is in PNG format and specified a length that is shorter than PNG header.
    PngShorterThanHeader,
    /// The enclosed PNG is not in RGBA, which is invalid: https://blogs.msdn.microsoft.com/oldnewthing/20101022-00/?p=12473/.
    PngNotRgba,

    /// The entry is in BMP format and specified a data size that is not correct for the image and optional mask data.
    InvalidDataSize,

    /// The dimensions specified by the entry does not match the dimensions in the header of the enclosed image.
    ImageEntryDimensionMismatch {
        /// The mismatched subimage's type
        format: IcoEntryImageFormat,
        /// The dimensions specified by the entry
        entry: (u16, u16),
        /// The dimensions of the image itself
        image: (u32, u32)
    },
}

impl fmt::Display for DecoderError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DecoderError::NoEntries =>
                f.write_str("ICO directory contains no image"),
            DecoderError::IcoEntryTooManyPlanesOrHotspot =>
                f.write_str("ICO image entry has too many color planes or too large hotspot value"),
            DecoderError::IcoEntryTooManyBitsPerPixelOrHotspot =>
                f.write_str("ICO image entry has too many bits per pixel or too large hotspot value"),
            DecoderError::PngShorterThanHeader =>
                f.write_str("Entry specified a length that is shorter than PNG header!"),
            DecoderError::PngNotRgba =>
                f.write_str("The PNG is not in RGBA format!"),
            DecoderError::InvalidDataSize =>
                f.write_str("ICO image data size did not match expected size"),
            DecoderError::ImageEntryDimensionMismatch { format, entry, image } =>
                f.write_fmt(format_args!("Entry{:?} and {}{:?} dimensions do not match!", entry, format, image)),
        }
    }
}

impl From<DecoderError> for ImageError {
    fn from(e: DecoderError) -> ImageError {
        ImageError::Decoding(DecodingError::new(ImageFormat::Ico.into(), e))
    }
}

impl error::Error for DecoderError {}

/// The image formats an ICO may contain
#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
enum IcoEntryImageFormat {
    /// PNG in ARGB
    Png,
    /// BMP with optional alpha mask
    Bmp,
}

impl fmt::Display for IcoEntryImageFormat {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            IcoEntryImageFormat::Png => "PNG",
            IcoEntryImageFormat::Bmp => "BMP",
        })
    }
}

impl Into<ImageFormat> for IcoEntryImageFormat {
    fn into(self) -> ImageFormat {
        match self {
            IcoEntryImageFormat::Png => ImageFormat::Png,
            IcoEntryImageFormat::Bmp => ImageFormat::Bmp,
        }
    }
}

/// An ico decoder
pub struct IcoDecoder<R: Read> {
    selected_entry: DirEntry,
    inner_decoder: InnerDecoder<R>,
}

enum InnerDecoder<R: Read> {
    BMP(BmpDecoder<R>),
    PNG(PngDecoder<R>),
}

#[derive(Clone, Copy, Default)]
struct DirEntry {
    width: u8,
    height: u8,
    // We ignore some header fields as they will be replicated in the PNG, BMP and they are not
    // necessary for determining the best_entry.
    #[allow(unused)]
    color_count: u8,
    // Wikipedia has this to say:
    // Although Microsoft's technical documentation states that this value must be zero, the icon
    // encoder built into .NET (System.Drawing.Icon.Save) sets this value to 255. It appears that
    // the operating system ignores this value altogether.
    #[allow(unused)]
    reserved: u8,

    // We ignore some header fields as they will be replicated in the PNG, BMP and they are not
    // necessary for determining the best_entry.
    #[allow(unused)]
    num_color_planes: u16,
    bits_per_pixel: u16,

    image_length: u32,
    image_offset: u32,
}

impl<R: Read + Seek> IcoDecoder<R> {
    /// Create a new decoder that decodes from the stream ```r```
    pub fn new(mut r: R) -> ImageResult<IcoDecoder<R>> {
        let entries = read_entries(&mut r)?;
        let entry = best_entry(entries)?;
        let decoder = entry.decoder(r)?;

        Ok(IcoDecoder {
            selected_entry: entry,
            inner_decoder: decoder,
        })
    }
}

fn read_entries<R: Read>(r: &mut R) -> ImageResult<Vec<DirEntry>> {
    let _reserved = r.read_u16::<LittleEndian>()?;
    let _type = r.read_u16::<LittleEndian>()?;
    let count = r.read_u16::<LittleEndian>()?;
    (0..count).map(|_| read_entry(r)).collect()
}

fn read_entry<R: Read>(r: &mut R) -> ImageResult<DirEntry> {
    Ok(DirEntry {
        width: r.read_u8()?,
        height: r.read_u8()?,
        color_count: r.read_u8()?,
        reserved: r.read_u8()?,
        num_color_planes: {
            // This may be either the number of color planes (0 or 1), or the horizontal coordinate
            // of the hotspot for CUR files.
            let num = r.read_u16::<LittleEndian>()?;
            if num > 256 {
                return Err(DecoderError::IcoEntryTooManyPlanesOrHotspot.into());
            }
            num
        },
        bits_per_pixel: {
            // This may be either the bit depth (may be 0 meaning unspecified),
            // or the vertical coordinate of the hotspot for CUR files.
            let num = r.read_u16::<LittleEndian>()?;
            if num > 256 {
                return Err(DecoderError::IcoEntryTooManyBitsPerPixelOrHotspot.into());
            }
            num
        },
        image_length: r.read_u32::<LittleEndian>()?,
        image_offset: r.read_u32::<LittleEndian>()?,
    })
}

/// Find the entry with the highest (color depth, size).
fn best_entry(mut entries: Vec<DirEntry>) -> ImageResult<DirEntry> {
    let mut best = entries.pop().ok_or(DecoderError::NoEntries)?;

    let mut best_score = (
        best.bits_per_pixel,
        u32::from(best.real_width()) * u32::from(best.real_height()),
    );

    for entry in entries {
        let score = (
            entry.bits_per_pixel,
            u32::from(entry.real_width()) * u32::from(entry.real_height()),
        );
        if score > best_score {
            best = entry;
            best_score = score;
        }
    }
    Ok(best)
}

impl DirEntry {
    fn real_width(&self) -> u16 {
        match self.width {
            0 => 256,
            w => u16::from(w),
        }
    }

    fn real_height(&self) -> u16 {
        match self.height {
            0 => 256,
            h => u16::from(h),
        }
    }

    fn matches_dimensions(&self, width: u32, height: u32) -> bool {
        u32::from(self.real_width()) == width && u32::from(self.real_height()) == height
    }

    fn seek_to_start<R: Read + Seek>(&self, r: &mut R) -> ImageResult<()> {
        r.seek(SeekFrom::Start(u64::from(self.image_offset)))?;
        Ok(())
    }

    fn is_png<R: Read + Seek>(&self, r: &mut R) -> ImageResult<bool> {
        self.seek_to_start(r)?;

        // Read the first 8 bytes to sniff the image.
        let mut signature = [0u8; 8];
        r.read_exact(&mut signature)?;

        Ok(signature == PNG_SIGNATURE)
    }

    fn decoder<R: Read + Seek>(&self, mut r: R) -> ImageResult<InnerDecoder<R>> {
        let is_png = self.is_png(&mut r)?;
        self.seek_to_start(&mut r)?;

        if is_png {
            Ok(PNG(PngDecoder::new(r)?))
        } else {
            Ok(BMP(BmpDecoder::new_with_ico_format(r)?))
        }
    }
}

/// Wrapper struct around a `Cursor<Vec<u8>>`
pub struct IcoReader<R>(Cursor<Vec<u8>>, PhantomData<R>);
impl<R> Read for IcoReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
        if self.0.position() == 0 && buf.is_empty() {
            mem::swap(buf, self.0.get_mut());
            Ok(buf.len())
        } else {
            self.0.read_to_end(buf)
        }
    }
}

impl<'a, R: 'a + Read + Seek> ImageDecoder<'a> for IcoDecoder<R> {
    type Reader = IcoReader<R>;

    fn dimensions(&self) -> (u32, u32) {
        match self.inner_decoder {
            BMP(ref decoder) => decoder.dimensions(),
            PNG(ref decoder) => decoder.dimensions(),
        }
    }

    fn color_type(&self) -> ColorType {
        match self.inner_decoder {
            BMP(ref decoder) => decoder.color_type(),
            PNG(ref decoder) => decoder.color_type(),
        }
    }

    fn into_reader(self) -> ImageResult<Self::Reader> {
        Ok(IcoReader(Cursor::new(image::decoder_to_vec(self)?), PhantomData))
    }

    fn read_image(self, buf: &mut [u8]) -> ImageResult<()> {
        assert_eq!(u64::try_from(buf.len()), Ok(self.total_bytes()));
        match self.inner_decoder {
            PNG(decoder) => {
                if self.selected_entry.image_length < PNG_SIGNATURE.len() as u32 {
                    return Err(DecoderError::PngShorterThanHeader.into());
                }

                // Check if the image dimensions match the ones in the image data.
                let (width, height) = decoder.dimensions();
                if !self.selected_entry.matches_dimensions(width, height) {
                    return Err(DecoderError::ImageEntryDimensionMismatch {
                        format: IcoEntryImageFormat::Png,
                        entry: (self.selected_entry.real_width(), self.selected_entry.real_height()),
                        image: (width, height)
                    }.into());
                }

                // Embedded PNG images can only be of the 32BPP RGBA format.
                // https://blogs.msdn.microsoft.com/oldnewthing/20101022-00/?p=12473/
                if decoder.color_type() != ColorType::Rgba8 {
                    return Err(DecoderError::PngNotRgba.into());
                }

                decoder.read_image(buf)
            }
            BMP(mut decoder) => {
                let (width, height) = decoder.dimensions();
                if !self.selected_entry.matches_dimensions(width, height) {
                    return Err(DecoderError::ImageEntryDimensionMismatch {
                        format: IcoEntryImageFormat::Bmp,
                        entry: (self.selected_entry.real_width(), self.selected_entry.real_height()),
                        image: (width, height)
                    }.into());
                }

                // The ICO decoder needs an alpha channel to apply the AND mask.
                if decoder.color_type() != ColorType::Rgba8 {
                    return Err(ImageError::Unsupported(UnsupportedError::from_format_and_kind(
                        ImageFormat::Bmp.into(),
                        UnsupportedErrorKind::Color(decoder.color_type().into()),
                    )));
                }

                decoder.read_image_data(buf)?;

                let r = decoder.reader();
                let image_end = r.seek(SeekFrom::Current(0))?;
                let data_end =
                    u64::from(self.selected_entry.image_offset + self.selected_entry.image_length);

                let mask_row_bytes = ((width + 31) / 32) * 4;
                let mask_length = u64::from(mask_row_bytes) * u64::from(height);

                // data_end should be image_end + the mask length (mask_row_bytes * height).
                // According to
                // https://devblogs.microsoft.com/oldnewthing/20101021-00/?p=12483
                // the mask is required, but according to Wikipedia
                // https://en.wikipedia.org/wiki/ICO_(file_format)
                // the mask is not required. Unfortunately, Wikipedia does not have a citation
                // for that claim, so we can't be sure which is correct.
                if data_end >= image_end + mask_length {
                    // If there's an AND mask following the image, read and apply it.
                    for y in 0..height {
                        let mut x = 0;
                        for _ in 0..mask_row_bytes {
                            // Apply the bits of each byte until we reach the end of the row.
                            let mask_byte = r.read_u8()?;
                            for bit in (0..8).rev() {
                                if x >= width {
                                    break;
                                }
                                if mask_byte & (1 << bit) != 0 {
                                    // Set alpha channel to transparent.
                                    buf[((height - y - 1) * width + x) as usize * 4 + 3] = 0;
                                }
                                x += 1;
                            }
                        }
                    }

                    Ok(())
                } else if data_end == image_end {
                    // accept images with no mask data
                    Ok(())
                } else {
                    Err(DecoderError::InvalidDataSize.into())
                }
            }
        }
    }
}
<commit_msg>Resolve an internal capitalization inconsistency<commit_after>use byteorder::{LittleEndian, ReadBytesExt};
use std::convert::TryFrom;
use std::io::{self, Cursor, Read, Seek, SeekFrom};
use std::marker::PhantomData;
use std::{error, fmt, mem};

use crate::color::ColorType;
use crate::error::{DecodingError, ImageError, ImageResult, UnsupportedError, UnsupportedErrorKind};
use crate::image::{self, ImageDecoder, ImageFormat};

use self::InnerDecoder::*;
use crate::bmp::BmpDecoder;
use crate::png::PngDecoder;

// http://www.w3.org/TR/PNG-Structure.html
// The first eight bytes of a PNG file always contain the following (decimal) values:
const PNG_SIGNATURE: [u8; 8] = [137, 80, 78, 71, 13, 10, 26, 10];

/// Errors that can occur during decoding and parsing an ICO image or one of its enclosed images.
#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
enum DecoderError {
    /// The ICO directory is empty
    NoEntries,
    /// The number of color planes (0 or 1), or the horizontal coordinate of the hotspot for CUR files too big.
    IcoEntryTooManyPlanesOrHotspot,
    /// The bit depth (may be 0 meaning unspecified), or the vertical coordinate of the hotspot for CUR files too big.
    IcoEntryTooManyBitsPerPixelOrHotspot,

    /// The entry is in PNG format and specified a length that is shorter than PNG header.
    PngShorterThanHeader,
    /// The enclosed PNG is not in RGBA, which is invalid: https://blogs.msdn.microsoft.com/oldnewthing/20101022-00/?p=12473/.
    PngNotRgba,

    /// The entry is in BMP format and specified a data size that is not correct for the image and optional mask data.
    InvalidDataSize,

    /// The dimensions specified by the entry does not match the dimensions in the header of the enclosed image.
    ImageEntryDimensionMismatch {
        /// The mismatched subimage's type
        format: IcoEntryImageFormat,
        /// The dimensions specified by the entry
        entry: (u16, u16),
        /// The dimensions of the image itself
        image: (u32, u32)
    },
}

impl fmt::Display for DecoderError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DecoderError::NoEntries =>
                f.write_str("ICO directory contains no image"),
            DecoderError::IcoEntryTooManyPlanesOrHotspot =>
                f.write_str("ICO image entry has too many color planes or too large hotspot value"),
            DecoderError::IcoEntryTooManyBitsPerPixelOrHotspot =>
                f.write_str("ICO image entry has too many bits per pixel or too large hotspot value"),
            DecoderError::PngShorterThanHeader =>
                f.write_str("Entry specified a length that is shorter than PNG header!"),
            DecoderError::PngNotRgba =>
                f.write_str("The PNG is not in RGBA format!"),
            DecoderError::InvalidDataSize =>
                f.write_str("ICO image data size did not match expected size"),
            DecoderError::ImageEntryDimensionMismatch { format, entry, image } =>
                f.write_fmt(format_args!("Entry{:?} and {}{:?} dimensions do not match!", entry, format, image)),
        }
    }
}

impl From<DecoderError> for ImageError {
    fn from(e: DecoderError) -> ImageError {
        ImageError::Decoding(DecodingError::new(ImageFormat::Ico.into(), e))
    }
}

impl error::Error for DecoderError {}

/// The image formats an ICO may contain
#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
enum IcoEntryImageFormat {
    /// PNG in ARGB
    Png,
    /// BMP with optional alpha mask
    Bmp,
}

impl fmt::Display for IcoEntryImageFormat {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            IcoEntryImageFormat::Png => "PNG",
            IcoEntryImageFormat::Bmp => "BMP",
        })
    }
}

impl Into<ImageFormat> for IcoEntryImageFormat {
    fn into(self) -> ImageFormat {
        match self {
            IcoEntryImageFormat::Png => ImageFormat::Png,
            IcoEntryImageFormat::Bmp => ImageFormat::Bmp,
        }
    }
}

/// An ico decoder
pub struct IcoDecoder<R: Read> {
    selected_entry: DirEntry,
    inner_decoder: InnerDecoder<R>,
}

enum InnerDecoder<R: Read> {
    Bmp(BmpDecoder<R>),
    Png(PngDecoder<R>),
}

#[derive(Clone, Copy, Default)]
struct DirEntry {
    width: u8,
    height: u8,
    // We ignore some header fields as they will be replicated in the PNG, BMP and they are not
    // necessary for determining the best_entry.
    #[allow(unused)]
    color_count: u8,
    // Wikipedia has this to say:
    // Although Microsoft's technical documentation states that this value must be zero, the icon
    // encoder built into .NET (System.Drawing.Icon.Save) sets this value to 255. It appears that
    // the operating system ignores this value altogether.
    #[allow(unused)]
    reserved: u8,

    // We ignore some header fields as they will be replicated in the PNG, BMP and they are not
    // necessary for determining the best_entry.
    #[allow(unused)]
    num_color_planes: u16,
    bits_per_pixel: u16,

    image_length: u32,
    image_offset: u32,
}

impl<R: Read + Seek> IcoDecoder<R> {
    /// Create a new decoder that decodes from the stream ```r```
    pub fn new(mut r: R) -> ImageResult<IcoDecoder<R>> {
        let entries = read_entries(&mut r)?;
        let entry = best_entry(entries)?;
        let decoder = entry.decoder(r)?;

        Ok(IcoDecoder {
            selected_entry: entry,
            inner_decoder: decoder,
        })
    }
}

fn read_entries<R: Read>(r: &mut R) -> ImageResult<Vec<DirEntry>> {
    let _reserved = r.read_u16::<LittleEndian>()?;
    let _type = r.read_u16::<LittleEndian>()?;
    let count = r.read_u16::<LittleEndian>()?;
    (0..count).map(|_| read_entry(r)).collect()
}

fn read_entry<R: Read>(r: &mut R) -> ImageResult<DirEntry> {
    Ok(DirEntry {
        width: r.read_u8()?,
        height: r.read_u8()?,
        color_count: r.read_u8()?,
        reserved: r.read_u8()?,
        num_color_planes: {
            // This may be either the number of color planes (0 or 1), or the horizontal coordinate
            // of the hotspot for CUR files.
            let num = r.read_u16::<LittleEndian>()?;
            if num > 256 {
                return Err(DecoderError::IcoEntryTooManyPlanesOrHotspot.into());
            }
            num
        },
        bits_per_pixel: {
            // This may be either the bit depth (may be 0 meaning unspecified),
            // or the vertical coordinate of the hotspot for CUR files.
            let num = r.read_u16::<LittleEndian>()?;
            if num > 256 {
                return Err(DecoderError::IcoEntryTooManyBitsPerPixelOrHotspot.into());
            }
            num
        },
        image_length: r.read_u32::<LittleEndian>()?,
        image_offset: r.read_u32::<LittleEndian>()?,
    })
}

/// Find the entry with the highest (color depth, size).
fn best_entry(mut entries: Vec<DirEntry>) -> ImageResult<DirEntry> {
    let mut best = entries.pop().ok_or(DecoderError::NoEntries)?;

    let mut best_score = (
        best.bits_per_pixel,
        u32::from(best.real_width()) * u32::from(best.real_height()),
    );

    for entry in entries {
        let score = (
            entry.bits_per_pixel,
            u32::from(entry.real_width()) * u32::from(entry.real_height()),
        );
        if score > best_score {
            best = entry;
            best_score = score;
        }
    }
    Ok(best)
}

impl DirEntry {
    fn real_width(&self) -> u16 {
        match self.width {
            0 => 256,
            w => u16::from(w),
        }
    }

    fn real_height(&self) -> u16 {
        match self.height {
            0 => 256,
            h => u16::from(h),
        }
    }

    fn matches_dimensions(&self, width: u32, height: u32) -> bool {
        u32::from(self.real_width()) == width && u32::from(self.real_height()) == height
    }

    fn seek_to_start<R: Read + Seek>(&self, r: &mut R) -> ImageResult<()> {
        r.seek(SeekFrom::Start(u64::from(self.image_offset)))?;
        Ok(())
    }

    fn is_png<R: Read + Seek>(&self, r: &mut R) -> ImageResult<bool> {
        self.seek_to_start(r)?;

        // Read the first 8 bytes to sniff the image.
        let mut signature = [0u8; 8];
        r.read_exact(&mut signature)?;

        Ok(signature == PNG_SIGNATURE)
    }

    fn decoder<R: Read + Seek>(&self, mut r: R) -> ImageResult<InnerDecoder<R>> {
        let is_png = self.is_png(&mut r)?;
        self.seek_to_start(&mut r)?;

        if is_png {
            Ok(Png(PngDecoder::new(r)?))
        } else {
            Ok(Bmp(BmpDecoder::new_with_ico_format(r)?))
        }
    }
}

/// Wrapper struct around a `Cursor<Vec<u8>>`
pub struct IcoReader<R>(Cursor<Vec<u8>>, PhantomData<R>);
impl<R> Read for IcoReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
        if self.0.position() == 0 && buf.is_empty() {
            mem::swap(buf, self.0.get_mut());
            Ok(buf.len())
        } else {
            self.0.read_to_end(buf)
        }
    }
}

impl<'a, R: 'a + Read + Seek> ImageDecoder<'a> for IcoDecoder<R> {
    type Reader = IcoReader<R>;

    fn dimensions(&self) -> (u32, u32) {
        match self.inner_decoder {
            Bmp(ref decoder) => decoder.dimensions(),
            Png(ref decoder) => decoder.dimensions(),
        }
    }

    fn color_type(&self) -> ColorType {
        match self.inner_decoder {
            Bmp(ref decoder) => decoder.color_type(),
            Png(ref decoder) => decoder.color_type(),
        }
    }

    fn into_reader(self) -> ImageResult<Self::Reader> {
        Ok(IcoReader(Cursor::new(image::decoder_to_vec(self)?), PhantomData))
    }

    fn read_image(self, buf: &mut [u8]) -> ImageResult<()> {
        assert_eq!(u64::try_from(buf.len()), Ok(self.total_bytes()));
        match self.inner_decoder {
            Png(decoder) => {
                if self.selected_entry.image_length < PNG_SIGNATURE.len() as u32 {
                    return Err(DecoderError::PngShorterThanHeader.into());
                }

                // Check if the image dimensions match the ones in the image data.
                let (width, height) = decoder.dimensions();
                if !self.selected_entry.matches_dimensions(width, height) {
                    return Err(DecoderError::ImageEntryDimensionMismatch {
                        format: IcoEntryImageFormat::Png,
                        entry: (self.selected_entry.real_width(), self.selected_entry.real_height()),
                        image: (width, height)
                    }.into());
                }

                // Embedded PNG images can only be of the 32BPP RGBA format.
                // https://blogs.msdn.microsoft.com/oldnewthing/20101022-00/?p=12473/
                if decoder.color_type() != ColorType::Rgba8 {
                    return Err(DecoderError::PngNotRgba.into());
                }

                decoder.read_image(buf)
            }
            Bmp(mut decoder) => {
                let (width, height) = decoder.dimensions();
                if !self.selected_entry.matches_dimensions(width, height) {
                    return Err(DecoderError::ImageEntryDimensionMismatch {
                        format: IcoEntryImageFormat::Bmp,
                        entry: (self.selected_entry.real_width(), self.selected_entry.real_height()),
                        image: (width, height)
                    }.into());
                }

                // The ICO decoder needs an alpha channel to apply the AND mask.
                if decoder.color_type() != ColorType::Rgba8 {
                    return Err(ImageError::Unsupported(UnsupportedError::from_format_and_kind(
                        ImageFormat::Bmp.into(),
                        UnsupportedErrorKind::Color(decoder.color_type().into()),
                    )));
                }

                decoder.read_image_data(buf)?;

                let r = decoder.reader();
                let image_end = r.seek(SeekFrom::Current(0))?;
                let data_end =
                    u64::from(self.selected_entry.image_offset + self.selected_entry.image_length);

                let mask_row_bytes = ((width + 31) / 32) * 4;
                let mask_length = u64::from(mask_row_bytes) * u64::from(height);

                // data_end should be image_end + the mask length (mask_row_bytes * height).
                // According to
                // https://devblogs.microsoft.com/oldnewthing/20101021-00/?p=12483
                // the mask is required, but according to Wikipedia
                // https://en.wikipedia.org/wiki/ICO_(file_format)
                // the mask is not required. Unfortunately, Wikipedia does not have a citation
                // for that claim, so we can't be sure which is correct.
                if data_end >= image_end + mask_length {
                    // If there's an AND mask following the image, read and apply it.
                    for y in 0..height {
                        let mut x = 0;
                        for _ in 0..mask_row_bytes {
                            // Apply the bits of each byte until we reach the end of the row.
                            let mask_byte = r.read_u8()?;
                            for bit in (0..8).rev() {
                                if x >= width {
                                    break;
                                }
                                if mask_byte & (1 << bit) != 0 {
                                    // Set alpha channel to transparent.
                                    buf[((height - y - 1) * width + x) as usize * 4 + 3] = 0;
                                }
                                x += 1;
                            }
                        }
                    }

                    Ok(())
                } else if data_end == image_end {
                    // accept images with no mask data
                    Ok(())
                } else {
                    Err(DecoderError::InvalidDataSize.into())
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use super::{Result, Specifier, SpecifierClass, SpecifierStack};
use std::rc::Rc;
use std::str::FromStr;

pub fn spec(s: &str) -> Result<Rc<Specifier>> {
    Specifier::from_stack(&SpecifierStack::from_str(s)?)
}

fn some_checks(s: &str) -> Result<()> {
    #[cfg(not(feature = "ssl"))]
    {
        if s.starts_with("wss://") {
            Err("SSL is not compiled in. Use ws:// or get/make another Websocat build.\nYou can also try to workaround missing SSL by using ws-c:cmd:socat trick (see some ws-c: example)")?
        }
    }

    #[cfg(not(any(target_os = "linux", target_os = "android")))]
    {
        if s.starts_with("abstract") {
            warn!("Abstract-namespaced UNIX sockets are unlikely to be supported here");
        }
    }

    if s.starts_with("open:") {
        return Err("There is no `open:` address type. Consider `open-async:` or `readfile:` or `writefile:` or `appendfile:`")?;
    }

    #[cfg(not(unix))]
    {
        if s.starts_with("unix") || s.starts_with("abstract") {
            Err("`unix*:` or `abstract*:` are not supported in this Websocat build")?
        }
    }

    #[cfg(not(feature = "tokio-process"))]
    {
        if s.starts_with("sh-c:") {
            Err("`sh-c:` is not supported in this Websocat build")?
        } else if s.starts_with("exec:") {
            Err("`exec:` is not supported in this Websocat build")?
        }
    }

    Ok(())
}

impl FromStr for SpecifierStack {
    type Err = Box<::std::error::Error>;
    #[cfg_attr(feature = "cargo-clippy", allow(cyclomatic_complexity))]
    fn from_str(s: &str) -> Result<SpecifierStack> {
        some_checks(s)?;

        let mut s = s.to_string();
        let mut overlays = vec![];
        let addrtype;
        let addr;
        let mut found = false;

        'a: loop {
            macro_rules! my {
                ($x:expr) => {
                    for pre in $x.get_prefixes() {
                        if s.starts_with(pre) {
                            let rest = &s[pre.len()..].to_string();
                            if let Some(a) = $x.alias_info() {
                                s = format!("{}{}", a, rest);
                                continue 'a;
                            } else if $x.is_overlay() {
                                overlays.push(Rc::new($x) as Rc<SpecifierClass>);
                                s = rest.to_string();
                                continue 'a;
                            } else {
                                addr = rest.to_string();
                                addrtype = Rc::new($x) as Rc<SpecifierClass>;
                                #[allow(unused_assignments)]
                                {
                                    found = true;
                                }
                                break 'a;
                            }
                        }
                    }
                };
            }
            list_of_all_specifier_classes!(my);
            if !found {
                Err(format!("Unknown address or overlay type of `{}`", s))?;
            }
        }

        Ok(SpecifierStack {
            addr,
            addrtype,
            overlays,
        })
    }
}

impl Specifier {
    pub fn from_stack(st: &SpecifierStack) -> Result<Rc<Specifier>> {
        let mut x = st.addrtype.construct(st.addr.as_str())?;
        for overlay in st.overlays.iter().rev() {
            x = overlay.construct_overlay(x)?;
        }
        Ok(x)
    }
}
<commit_msg>Better error message for unknown specifiers<commit_after>use super::{Result, Specifier, SpecifierClass, SpecifierStack};
use std::rc::Rc;
use std::str::FromStr;

pub fn spec(s: &str) -> Result<Rc<Specifier>> {
    Specifier::from_stack(&SpecifierStack::from_str(s)?)
}

fn some_checks(s: &str) -> Result<()> {
    #[cfg(not(feature = "ssl"))]
    {
        if s.starts_with("wss://") {
            Err("SSL is not compiled in. Use ws:// or get/make another Websocat build.\nYou can also try to workaround missing SSL by using ws-c:cmd:socat trick (see some ws-c: example)")?
        }
    }

    #[cfg(not(any(target_os = "linux", target_os = "android")))]
    {
        if s.starts_with("abstract") {
            warn!("Abstract-namespaced UNIX sockets are unlikely to be supported here");
        }
    }

    if s.starts_with("open:") {
        return Err("There is no `open:` address type. Consider `open-async:` or `readfile:` or `writefile:` or `appendfile:`")?;
    }

    #[cfg(not(unix))]
    {
        if s.starts_with("unix") || s.starts_with("abstract") {
            Err("`unix*:` or `abstract*:` are not supported in this Websocat build")?
        }
    }

    #[cfg(not(feature = "tokio-process"))]
    {
        if s.starts_with("sh-c:") {
            Err("`sh-c:` is not supported in this Websocat build")?
        } else if s.starts_with("exec:") {
            Err("`exec:` is not supported in this Websocat build")?
        }
    }

    Ok(())
}

impl FromStr for SpecifierStack {
    type Err = Box<::std::error::Error>;
    #[cfg_attr(feature = "cargo-clippy", allow(cyclomatic_complexity))]
    fn from_str(s: &str) -> Result<SpecifierStack> {
        some_checks(s)?;

        let mut s = s.to_string();
        let mut overlays = vec![];
        let addrtype;
        let addr;
        let mut found = false;

        'a: loop {
            macro_rules! my {
                ($x:expr) => {
                    for pre in $x.get_prefixes() {
                        if s.starts_with(pre) {
                            let rest = &s[pre.len()..].to_string();
                            if let Some(a) = $x.alias_info() {
                                s = format!("{}{}", a, rest);
                                continue 'a;
                            } else if $x.is_overlay() {
                                overlays.push(Rc::new($x) as Rc<SpecifierClass>);
                                s = rest.to_string();
                                continue 'a;
                            } else {
                                addr = rest.to_string();
                                addrtype = Rc::new($x) as Rc<SpecifierClass>;
                                #[allow(unused_assignments)]
                                {
                                    found = true;
                                }
                                break 'a;
                            }
                        }
                    }
                };
            }
            list_of_all_specifier_classes!(my);
            if !found {
                if let Some(colon) = s.find(':') {
                    Err(format!("Unknown address or overlay type of `{}:`", &s[..colon]))?;
                } else {
                    Err(format!("Unknown address or overlay type of `{}`\nMaybe you forgot the `:` character?", s))?;
                }
            }
        }

        Ok(SpecifierStack {
            addr,
            addrtype,
            overlays,
        })
    }
}

impl Specifier {
    pub fn from_stack(st: &SpecifierStack) -> Result<Rc<Specifier>> {
        let mut x = st.addrtype.construct(st.addr.as_str())?;
        for overlay in st.overlays.iter().rev() {
            x = overlay.construct_overlay(x)?;
        }
        Ok(x)
    }
}
<|endoftext|>


<commit_before><commit_msg>Renaming<commit_after><|endoftext|>


<commit_before>use std::{collections::BTreeMap, io};

use rayon::prelude::*;

use self::reader::LogReader;

use super::*;

pub struct LogIter {
    pub config: Config,
    pub segment_iter: Box<dyn Iterator<Item = (Lsn, LogId)>>,
    pub segment_base: Option<LogId>,
    pub max_lsn: Lsn,
    pub cur_lsn: Lsn,
}

impl Iterator for LogIter {
    type Item = (LogKind, PageId, Lsn, DiskPtr, usize);

    fn next(&mut self) -> Option<Self::Item> {
        // If segment is None, get next on segment_iter, panic
        // if we can't read something we expect to be able to,
        // return None if there are no more remaining segments.
        loop {
            let remaining_seg_too_small_for_msg = !valid_entry_offset(
                self.cur_lsn as LogId,
                self.config.io_buf_size,
            );

            if self.segment_base.is_none() || remaining_seg_too_small_for_msg {
                if let Some((next_lsn, next_lid)) = self.segment_iter.next() {
                    assert!(
                        next_lsn + (self.config.io_buf_size as Lsn)
                            >= self.cur_lsn,
                        "caller is responsible for providing segments \
                         that contain the initial cur_lsn value or higher"
                    );

                    #[cfg(target_os = "linux")]
                    self.fadvise_willneed(next_lid);

                    if let Err(e) = self.read_segment(next_lsn, next_lid) {
                        debug!(
                            "hit snap while reading segments in \
                             iterator: {:?}",
                            e
                        );
                        return None;
                    }
                } else {
                    trace!("no segments remaining to iterate over");
                    return None;
                }
            }

            if self.cur_lsn > self.max_lsn {
                // all done
                trace!("hit max_lsn {} in iterator, stopping", self.max_lsn);
                return None;
            }

            let lid = self.segment_base.unwrap()
                + (self.cur_lsn % self.config.io_buf_size as Lsn) as LogId;

            let f = &self.config.file;

            match f.read_message(lid, self.cur_lsn, &self.config) {
                Ok(LogRead::Blob(header, _buf, blob_ptr)) => {
                    trace!("read blob flush in LogIter::next");
                    let sz = MSG_HEADER_LEN + BLOB_INLINE_LEN;
                    self.cur_lsn += sz as Lsn;

                    return Some((
                        LogKind::from(header.kind),
                        header.pid,
                        header.lsn,
                        DiskPtr::Blob(lid, blob_ptr),
                        sz,
                    ));
                }
                Ok(LogRead::Inline(header, _buf, on_disk_len)) => {
                    trace!("read inline flush in LogIter::next");
                    let sz = MSG_HEADER_LEN + on_disk_len as usize;
                    self.cur_lsn += sz as Lsn;

                    return Some((
                        LogKind::from(header.kind),
                        header.pid,
                        header.lsn,
                        DiskPtr::Inline(lid),
                        sz,
                    ));
                }
                Ok(LogRead::BatchManifest(last_lsn_in_batch)) => {
                    if last_lsn_in_batch > self.max_lsn {
                        return None;
                    } else {
                        self.cur_lsn +=
                            (MSG_HEADER_LEN + BATCH_MANIFEST_INLINE_LEN) as Lsn;
                        continue;
                    }
                }
                Ok(LogRead::Failed(_, on_disk_len)) => {
                    trace!("read zeroed in LogIter::next");
                    self.cur_lsn +=
                        Lsn::from(MSG_HEADER_LEN as u32 + on_disk_len);
                }
                Ok(LogRead::Corrupted(_len)) => {
                    trace!(
                        "read corrupted msg in LogIter::next as lid {} lsn {}",
                        lid,
                        self.cur_lsn
                    );
                    return None;
                }
                Ok(LogRead::Pad(_lsn)) => {
                    self.segment_base.take();

                    continue;
                }
                Ok(LogRead::DanglingBlob(header, blob_ptr)) => {
                    debug!(
                        "encountered dangling blob \
                         pointer at lsn {} ptr {}",
                        header.lsn, blob_ptr
                    );
                    self.cur_lsn += (MSG_HEADER_LEN + BLOB_INLINE_LEN) as Lsn;
                    continue;
                }
                Err(e) => {
                    debug!(
                        "failed to read log message at lid {} \
                         with expected lsn {} during iteration: {}",
                        lid, self.cur_lsn, e
                    );
                    return None;
                }
            }
        }
    }
}

impl LogIter {
    /// read a segment of log messages. Only call after
    /// pausing segment rewriting on the segment accountant!
    fn read_segment(&mut self, lsn: Lsn, offset: LogId) -> Result<()> {
        trace!(
            "LogIter::read_segment lsn: {:?} cur_lsn: {:?}",
            lsn,
            self.cur_lsn
        );
        // we add segment_len to this check because we may be getting the
        // initial segment that is a bit behind where we left off before.
        assert!(lsn + self.config.io_buf_size as Lsn >= self.cur_lsn);
        let f = &self.config.file;
        let segment_header = f.read_segment_header(offset)?;
        if offset % self.config.io_buf_size as LogId != 0 {
            debug!("segment offset not divisible by segment length");
            return Err(Error::Corruption {
                at: DiskPtr::Inline(offset),
            });
        }
        if segment_header.lsn % self.config.io_buf_size as Lsn != 0 {
            debug!(
                "expected a segment header lsn that is divisible \
                 by the io_buf_size ({}) instead it was {}",
                self.config.io_buf_size, segment_header.lsn
            );
            return Err(Error::Corruption {
                at: DiskPtr::Inline(offset),
            });
        }

        if segment_header.lsn != lsn {
            // this page was torn, nothing to read
            debug!(
                "segment header lsn ({}) != expected lsn ({})",
                segment_header.lsn, lsn
            );
            return Err(io::Error::new(
                io::ErrorKind::Other,
                "encountered torn segment",
            )
            .into());
        }

        trace!("read segment header {:?}", segment_header);

        self.cur_lsn = segment_header.lsn + SEG_HEADER_LEN as Lsn;
        self.segment_base = Some(offset);

        Ok(())
    }

    #[cfg(target_os = "linux")]
    fn fadvise_willneed(&self, lid: LogId) {
        use std::os::unix::io::AsRawFd;

        let f = &self.config.file;
        let ret = unsafe {
            libc::posix_fadvise(
                f.as_raw_fd(),
                libc::off_t::try_from(lid).unwrap(),
                libc::off_t::try_from(self.config.io_buf_size).unwrap(),
                libc::POSIX_FADV_WILLNEED,
            )
        };
        if ret != 0 {
            panic!(
                "failed to call fadvise: {}",
                std::io::Error::from_raw_os_error(ret)
            );
        }
    }
}

fn valid_entry_offset(lid: LogId, segment_len: usize) -> bool {
    let seg_start = lid / segment_len as LogId * segment_len as LogId;

    let max_lid = seg_start + segment_len as LogId - MSG_HEADER_LEN as LogId;

    let min_lid = seg_start + SEG_HEADER_LEN as LogId;

    lid >= min_lid && lid <= max_lid
}

// Scan the log file if we don't know of any Lsn offsets yet,
// and recover the order of segments, and the highest Lsn.
fn scan_segment_lsns(
    min: Lsn,
    config: &Config,
) -> Result<(BTreeMap<Lsn, LogId>, Lsn)> {
    let segment_len = LogId::try_from(config.io_buf_size).unwrap();

    let f = &config.file;
    let file_len = f.metadata()?.len();
    let segments = (file_len / segment_len)
        + if file_len % segment_len < LogId::try_from(SEG_HEADER_LEN).unwrap() {
            0
        } else {
            1
        };
    trace!(
        "file len: {} segment len {} segments: {}",
        file_len,
        segment_len,
        segments
    );
    let headers: Vec<(LogId, SegmentHeader)> = (0..segments)
        .into_par_iter()
        .filter_map(|idx| {
            let base_lid = idx * segment_len;
            let segment = f.read_segment_header(base_lid).ok()?;
            trace!(
                "SA scanned header at lid {} during startup: {:?}",
                base_lid,
                segment
            );
            if segment.ok && segment.lsn >= min {
                assert_ne!(segment.lsn, Lsn::max_value());
                Some((base_lid, segment))
            } else {
                trace!(
                    "not using segment at lid {}, ok: {} lsn: {} min lsn: {}",
                    base_lid,
                    segment.ok,
                    segment.lsn,
                    min
                );
                None
            }
        })
        .collect();

    let mut ordering = BTreeMap::new();
    let mut max_header_stable_lsn = 0;

    for (lid, header) in headers {
        max_header_stable_lsn =
            std::cmp::max(header.max_stable_lsn, max_header_stable_lsn);

        let old = ordering.insert(header.lsn, lid);
        assert_eq!(
            old, None,
            "duplicate segment LSN {} detected at both {} and {}, \
             one should have been zeroed out during recovery",
            header.lsn, ordering[&header.lsn], lid
        );
    }

    debug!("ordering before clearing tears: {:?}", ordering);

    // Check that the segments above max_header_stable_lsn
    // properly link their previous segment pointers.
    let ordering =
        clean_tail_tears(max_header_stable_lsn, ordering, config, &f)?;

    Ok((ordering, max_header_stable_lsn))
}

// This ensures that the last <# io buffers> segments on
// disk connect via their previous segment pointers in
// the header. This is important because we expect that
// the last <# io buffers> segments will join up, and we
// never reuse buffers within this safety range.
fn clean_tail_tears(
    max_header_stable_lsn: Lsn,
    mut ordering: BTreeMap<Lsn, LogId>,
    config: &Config,
    f: &std::fs::File,
) -> Result<BTreeMap<Lsn, LogId>> {
    let io_buf_size = config.io_buf_size as Lsn;
    let lowest_lsn_in_tail: Lsn =
        ((max_header_stable_lsn / io_buf_size) - 1) * io_buf_size;

    let mut expected_present = lowest_lsn_in_tail;
    let mut missing_item_in_tail = None;

    let logical_tail = ordering
        .range(lowest_lsn_in_tail..)
        .map(|(lsn, lid)| (*lsn, *lid))
        .take_while(|(lsn, _lid)| {
            let matches = expected_present == *lsn;
            if !matches {
                debug!(
                    "failed to find expected segment \
                     at lsn {}, tear detected",
                    expected_present
                );
                missing_item_in_tail = Some(expected_present);
            }
            expected_present += io_buf_size;
            matches
        })
        .collect::<Vec<_>>();

    let iter = LogIter {
        config: config.clone(),
        segment_iter: Box::new(logical_tail.into_iter()),
        segment_base: None,
        max_lsn: missing_item_in_tail.unwrap_or(Lsn::max_value()),
        cur_lsn: 0,
    };

    let tip: (Lsn, LogId) = iter
        .max_by_key(|(_kind, _pid, lsn, _ptr, _sz)| *lsn)
        .map(|(_, _, lsn, ptr, _)| (lsn, ptr.lid()))
        .unwrap_or_else(|| {
            if max_header_stable_lsn > 0 {
                (lowest_lsn_in_tail, ordering[&lowest_lsn_in_tail])
            } else {
                (0, 0)
            }
        });

    debug!(
        "filtering out segments after detected tear at lsn {} lid {}",
        tip.0, tip.1
    );
    for (lsn, lid) in ordering
        .range((std::ops::Bound::Excluded(tip.0), std::ops::Bound::Unbounded))
    {
        debug!("zeroing torn segment with lsn {} at lid {}", lsn, lid);

        f.pwrite_all(
            &*vec![MessageKind::Corrupted.into(); SEG_HEADER_LEN],
            *lid,
        )?;
        if !config.temporary {
            f.sync_all()?;
        }
    }

    ordering = ordering
        .into_iter()
        .filter(|&(lsn, _lid)| lsn <= tip.0)
        .collect();

    Ok(ordering)
}

pub(super) fn raw_segment_iter_from(
    lsn: Lsn,
    config: &Config,
) -> Result<(LogIter, Lsn)> {
    let segment_len = config.io_buf_size as Lsn;
    let normalized_lsn = lsn / segment_len * segment_len;

    let (ordering, max_header_stable_lsn) = scan_segment_lsns(0, &config)?;

    // find the last stable tip, to properly handle batch manifests.
    let tip_segment_iter =
        Box::new(ordering.iter().map(|(a, b)| (*a, *b)).last().into_iter());
    trace!(
        "trying to find the max stable tip for \
         bounding batch manifests with segment iter {:?} \
         of segments >= max_header_stable_lsn {}",
        tip_segment_iter,
        max_header_stable_lsn
    );

    let mut tip_iter = LogIter {
        config: config.clone(),
        max_lsn: Lsn::max_value(),
        cur_lsn: 0,
        segment_base: None,
        segment_iter: tip_segment_iter,
    };

    // run the iterator to the end so
    // we can grab its current lsn, inclusive
    // of any zeroed messages and other
    // legit items it may not have returned
    // in the actual iterator.
    while let Some(_) = tip_iter.next() {}

    let tip = tip_iter.cur_lsn;

    trace!("found max stable tip: {}", tip);

    trace!(
        "generated iterator over segments {:?} with lsn >= {}",
        ordering,
        normalized_lsn,
    );

    let segment_iter = Box::new(
        ordering
            .into_iter()
            .filter(move |&(l, _)| l >= normalized_lsn),
    );

    Ok((
        LogIter {
            config: config.clone(),
            max_lsn: tip,
            cur_lsn: 0,
            segment_base: None,
            segment_iter,
        },
        max_header_stable_lsn,
    ))
}
<commit_msg>gooder log message<commit_after>use std::{collections::BTreeMap, io};

use rayon::prelude::*;

use self::reader::LogReader;

use super::*;

pub struct LogIter {
    pub config: Config,
    pub segment_iter: Box<dyn Iterator<Item = (Lsn, LogId)>>,
    pub segment_base: Option<LogId>,
    pub max_lsn: Lsn,
    pub cur_lsn: Lsn,
}

impl Iterator for LogIter {
    type Item = (LogKind, PageId, Lsn, DiskPtr, usize);

    fn next(&mut self) -> Option<Self::Item> {
        // If segment is None, get next on segment_iter, panic
        // if we can't read something we expect to be able to,
        // return None if there are no more remaining segments.
        loop {
            let remaining_seg_too_small_for_msg = !valid_entry_offset(
                self.cur_lsn as LogId,
                self.config.io_buf_size,
            );

            if self.segment_base.is_none() || remaining_seg_too_small_for_msg {
                if let Some((next_lsn, next_lid)) = self.segment_iter.next() {
                    assert!(
                        next_lsn + (self.config.io_buf_size as Lsn)
                            >= self.cur_lsn,
                        "caller is responsible for providing segments \
                         that contain the initial cur_lsn value or higher"
                    );

                    #[cfg(target_os = "linux")]
                    self.fadvise_willneed(next_lid);

                    if let Err(e) = self.read_segment(next_lsn, next_lid) {
                        debug!(
                            "hit snap while reading segments in \
                             iterator: {:?}",
                            e
                        );
                        return None;
                    }
                } else {
                    trace!("no segments remaining to iterate over");
                    return None;
                }
            }

            if self.cur_lsn > self.max_lsn {
                // all done
                trace!("hit max_lsn {} in iterator, stopping", self.max_lsn);
                return None;
            }

            let lid = self.segment_base.unwrap()
                + (self.cur_lsn % self.config.io_buf_size as Lsn) as LogId;

            let f = &self.config.file;

            match f.read_message(lid, self.cur_lsn, &self.config) {
                Ok(LogRead::Blob(header, _buf, blob_ptr)) => {
                    trace!("read blob flush in LogIter::next");
                    let sz = MSG_HEADER_LEN + BLOB_INLINE_LEN;
                    self.cur_lsn += sz as Lsn;

                    return Some((
                        LogKind::from(header.kind),
                        header.pid,
                        header.lsn,
                        DiskPtr::Blob(lid, blob_ptr),
                        sz,
                    ));
                }
                Ok(LogRead::Inline(header, _buf, on_disk_len)) => {
                    trace!("read inline flush in LogIter::next");
                    let sz = MSG_HEADER_LEN + on_disk_len as usize;
                    self.cur_lsn += sz as Lsn;

                    return Some((
                        LogKind::from(header.kind),
                        header.pid,
                        header.lsn,
                        DiskPtr::Inline(lid),
                        sz,
                    ));
                }
                Ok(LogRead::BatchManifest(last_lsn_in_batch)) => {
                    if last_lsn_in_batch > self.max_lsn {
                        return None;
                    } else {
                        self.cur_lsn +=
                            (MSG_HEADER_LEN + BATCH_MANIFEST_INLINE_LEN) as Lsn;
                        continue;
                    }
                }
                Ok(LogRead::Failed(_, on_disk_len)) => {
                    trace!("read zeroed in LogIter::next");
                    self.cur_lsn +=
                        Lsn::from(MSG_HEADER_LEN as u32 + on_disk_len);
                }
                Ok(LogRead::Corrupted(_len)) => {
                    trace!(
                        "read corrupted msg in LogIter::next as lid {} lsn {}",
                        lid,
                        self.cur_lsn
                    );
                    return None;
                }
                Ok(LogRead::Pad(_lsn)) => {
                    self.segment_base.take();

                    continue;
                }
                Ok(LogRead::DanglingBlob(header, blob_ptr)) => {
                    debug!(
                        "encountered dangling blob \
                         pointer at lsn {} ptr {}",
                        header.lsn, blob_ptr
                    );
                    self.cur_lsn += (MSG_HEADER_LEN + BLOB_INLINE_LEN) as Lsn;
                    continue;
                }
                Err(e) => {
                    debug!(
                        "failed to read log message at lid {} \
                         with expected lsn {} during iteration: {}",
                        lid, self.cur_lsn, e
                    );
                    return None;
                }
            }
        }
    }
}

impl LogIter {
    /// read a segment of log messages. Only call after
    /// pausing segment rewriting on the segment accountant!
    fn read_segment(&mut self, lsn: Lsn, offset: LogId) -> Result<()> {
        trace!(
            "LogIter::read_segment lsn: {:?} cur_lsn: {:?}",
            lsn,
            self.cur_lsn
        );
        // we add segment_len to this check because we may be getting the
        // initial segment that is a bit behind where we left off before.
        assert!(lsn + self.config.io_buf_size as Lsn >= self.cur_lsn);
        let f = &self.config.file;
        let segment_header = f.read_segment_header(offset)?;
        if offset % self.config.io_buf_size as LogId != 0 {
            debug!("segment offset not divisible by segment length");
            return Err(Error::Corruption {
                at: DiskPtr::Inline(offset),
            });
        }
        if segment_header.lsn % self.config.io_buf_size as Lsn != 0 {
            debug!(
                "expected a segment header lsn that is divisible \
                 by the io_buf_size ({}) instead it was {}",
                self.config.io_buf_size, segment_header.lsn
            );
            return Err(Error::Corruption {
                at: DiskPtr::Inline(offset),
            });
        }

        if segment_header.lsn != lsn {
            // this page was torn, nothing to read
            debug!(
                "segment header lsn ({}) != expected lsn ({})",
                segment_header.lsn, lsn
            );
            return Err(io::Error::new(
                io::ErrorKind::Other,
                "encountered torn segment",
            )
            .into());
        }

        trace!("read segment header {:?}", segment_header);

        self.cur_lsn = segment_header.lsn + SEG_HEADER_LEN as Lsn;
        self.segment_base = Some(offset);

        Ok(())
    }

    #[cfg(target_os = "linux")]
    fn fadvise_willneed(&self, lid: LogId) {
        use std::os::unix::io::AsRawFd;

        let f = &self.config.file;
        let ret = unsafe {
            libc::posix_fadvise(
                f.as_raw_fd(),
                libc::off_t::try_from(lid).unwrap(),
                libc::off_t::try_from(self.config.io_buf_size).unwrap(),
                libc::POSIX_FADV_WILLNEED,
            )
        };
        if ret != 0 {
            panic!(
                "failed to call fadvise: {}",
                std::io::Error::from_raw_os_error(ret)
            );
        }
    }
}

fn valid_entry_offset(lid: LogId, segment_len: usize) -> bool {
    let seg_start = lid / segment_len as LogId * segment_len as LogId;

    let max_lid = seg_start + segment_len as LogId - MSG_HEADER_LEN as LogId;

    let min_lid = seg_start + SEG_HEADER_LEN as LogId;

    lid >= min_lid && lid <= max_lid
}

// Scan the log file if we don't know of any Lsn offsets yet,
// and recover the order of segments, and the highest Lsn.
fn scan_segment_lsns(
    min: Lsn,
    config: &Config,
) -> Result<(BTreeMap<Lsn, LogId>, Lsn)> {
    let segment_len = LogId::try_from(config.io_buf_size).unwrap();

    let f = &config.file;
    let file_len = f.metadata()?.len();
    let segments = (file_len / segment_len)
        + if file_len % segment_len < LogId::try_from(SEG_HEADER_LEN).unwrap() {
            0
        } else {
            1
        };
    trace!(
        "file len: {} segment len {} segments: {}",
        file_len,
        segment_len,
        segments
    );
    let headers: Vec<(LogId, SegmentHeader)> = (0..segments)
        .into_par_iter()
        .filter_map(|idx| {
            let base_lid = idx * segment_len;
            let segment = f.read_segment_header(base_lid).ok()?;
            trace!(
                "SA scanned header at lid {} during startup: {:?}",
                base_lid,
                segment
            );
            if segment.ok && segment.lsn >= min {
                assert_ne!(segment.lsn, Lsn::max_value());
                Some((base_lid, segment))
            } else {
                trace!(
                    "not using segment at lid {}, ok: {} lsn: {} min lsn: {}",
                    base_lid,
                    segment.ok,
                    segment.lsn,
                    min
                );
                None
            }
        })
        .collect();

    let mut ordering = BTreeMap::new();
    let mut max_header_stable_lsn = 0;

    for (lid, header) in headers {
        max_header_stable_lsn =
            std::cmp::max(header.max_stable_lsn, max_header_stable_lsn);

        let old = ordering.insert(header.lsn, lid);
        assert_eq!(
            old, None,
            "duplicate segment LSN {} detected at both {} and {}, \
             one should have been zeroed out during recovery",
            header.lsn, ordering[&header.lsn], lid
        );
    }

    debug!("ordering before clearing tears: {:?}", ordering);

    // Check that the segments above max_header_stable_lsn
    // properly link their previous segment pointers.
    let ordering =
        clean_tail_tears(max_header_stable_lsn, ordering, config, &f)?;

    Ok((ordering, max_header_stable_lsn))
}

// This ensures that the last <# io buffers> segments on
// disk connect via their previous segment pointers in
// the header. This is important because we expect that
// the last <# io buffers> segments will join up, and we
// never reuse buffers within this safety range.
fn clean_tail_tears(
    max_header_stable_lsn: Lsn,
    mut ordering: BTreeMap<Lsn, LogId>,
    config: &Config,
    f: &std::fs::File,
) -> Result<BTreeMap<Lsn, LogId>> {
    let io_buf_size = config.io_buf_size as Lsn;
    let lowest_lsn_in_tail: Lsn =
        ((max_header_stable_lsn / io_buf_size) - 1) * io_buf_size;

    let mut expected_present = lowest_lsn_in_tail;
    let mut missing_item_in_tail = None;

    let logical_tail = ordering
        .range(lowest_lsn_in_tail..)
        .map(|(lsn, lid)| (*lsn, *lid))
        .take_while(|(lsn, _lid)| {
            let matches = expected_present == *lsn;
            if !matches {
                debug!(
                    "failed to find expected segment \
                     at lsn {}, tear detected",
                    expected_present
                );
                missing_item_in_tail = Some(expected_present);
            }
            expected_present += io_buf_size;
            matches
        })
        .collect::<Vec<_>>();

    debug!(
        "in clean_tail_tears, found missing item in tail: {:?} \
         and we'll scan segments {:?} above lowest lsn {}",
        missing_item_in_tail, logical_tail, lowest_lsn_in_tail
    );

    let iter = LogIter {
        config: config.clone(),
        segment_iter: Box::new(logical_tail.into_iter()),
        segment_base: None,
        max_lsn: missing_item_in_tail.unwrap_or(Lsn::max_value()),
        cur_lsn: 0,
    };

    let tip: (Lsn, LogId) = iter
        .max_by_key(|(_kind, _pid, lsn, _ptr, _sz)| *lsn)
        .map(|(_, _, lsn, ptr, _)| (lsn, ptr.lid()))
        .unwrap_or_else(|| {
            if max_header_stable_lsn > 0 {
                (lowest_lsn_in_tail, ordering[&lowest_lsn_in_tail])
            } else {
                (0, 0)
            }
        });

    debug!(
        "filtering out segments after detected tear at lsn {} lid {}",
        tip.0, tip.1
    );
    for (lsn, lid) in ordering
        .range((std::ops::Bound::Excluded(tip.0), std::ops::Bound::Unbounded))
    {
        debug!("zeroing torn segment with lsn {} at lid {}", lsn, lid);

        f.pwrite_all(
            &*vec![MessageKind::Corrupted.into(); SEG_HEADER_LEN],
            *lid,
        )?;
        if !config.temporary {
            f.sync_all()?;
        }
    }

    ordering = ordering
        .into_iter()
        .filter(|&(lsn, _lid)| lsn <= tip.0)
        .collect();

    Ok(ordering)
}

pub(super) fn raw_segment_iter_from(
    lsn: Lsn,
    config: &Config,
) -> Result<(LogIter, Lsn)> {
    let segment_len = config.io_buf_size as Lsn;
    let normalized_lsn = lsn / segment_len * segment_len;

    let (ordering, max_header_stable_lsn) = scan_segment_lsns(0, &config)?;

    // find the last stable tip, to properly handle batch manifests.
    let tip_segment_iter =
        Box::new(ordering.iter().map(|(a, b)| (*a, *b)).last().into_iter());
    trace!(
        "trying to find the max stable tip for \
         bounding batch manifests with segment iter {:?} \
         of segments >= max_header_stable_lsn {}",
        tip_segment_iter,
        max_header_stable_lsn
    );

    let mut tip_iter = LogIter {
        config: config.clone(),
        max_lsn: Lsn::max_value(),
        cur_lsn: 0,
        segment_base: None,
        segment_iter: tip_segment_iter,
    };

    // run the iterator to the end so
    // we can grab its current lsn, inclusive
    // of any zeroed messages and other
    // legit items it may not have returned
    // in the actual iterator.
    while let Some(_) = tip_iter.next() {}

    let tip = tip_iter.cur_lsn;

    trace!("found max stable tip: {}", tip);

    trace!(
        "generated iterator over segments {:?} with lsn >= {}",
        ordering,
        normalized_lsn,
    );

    let segment_iter = Box::new(
        ordering
            .into_iter()
            .filter(move |&(l, _)| l >= normalized_lsn),
    );

    Ok((
        LogIter {
            config: config.clone(),
            max_lsn: tip,
            cur_lsn: 0,
            segment_base: None,
            segment_iter,
        },
        max_header_stable_lsn,
    ))
}
<|endoftext|>


<commit_before><commit_msg>Add "performance test".<commit_after><|endoftext|>


<commit_before>mod mb_calc {

    struct Complex {
        real: f64,
        imaginary: f64,
    }

    impl Complex {
        fn new(real: f64, imaginary: f64) -> Complex {
            Complex {
                real: real,
                imaginary: imaginary,
            }
        }
        fn abs(&self) -> f64 {
            let real = self.real * self.real;
            let imaginary = self.imaginary * self.imaginary;

            let sum = real + imaginary;

            let result = sum.sqrt();
            return result;
        }
    }

    #[test]
    fn new_gives_complex() {
        let c = Complex::new(2.0, 2.0 );
        assert_eq!(2.0 ,c.real);
        assert_eq!(2.0 ,c.imaginary);
    }

    #[test]
    fn abs_gives_absolute() {
        let res = 2.0 as f64;
        let c = Complex::new(1.0, 1.0);
        assert_eq!(res.sqrt(), c.abs());
    }
}
<commit_msg>Put complex into mb_calc::num<commit_after>mod mb_calc {
    mod num {
        struct Complex {
            real: f64,
            imaginary: f64,
        }

        impl Complex {
            fn new(real: f64, imaginary: f64) -> Complex {
                Complex {
                    real: real,
                    imaginary: imaginary,
                }
            }
            fn abs(&self) -> f64 {
                let real = self.real * self.real;
                let imaginary = self.imaginary * self.imaginary;

                let sum = real + imaginary;

                let result = sum.sqrt();
                return result;
            }
        }

        #[test]
        fn new_gives_complex() {
            let c = Complex::new(2.0, 2.0 );
            assert_eq!(2.0 ,c.real);
            assert_eq!(2.0 ,c.imaginary);
        }

        #[test]
        fn abs_gives_absolute() {
            let res = 2.0 as f64;
            let c = Complex::new(1.0, 1.0);
            assert_eq!(res.sqrt(), c.abs());
        }
    }
}
<|endoftext|>


<commit_before>extern crate crossterm;
extern crate tui;

use std::error::Error;
use std::io;

use tui::backend::CrosstermBackend;
use tui::layout::{Constraint, Direction, Layout};
use tui::style::{Color, Modifier, Style};
use tui::widgets::{Block, Borders, Paragraph, Text, Widget};
use tui::Terminal;

fn main() {
    let mut terminal = Terminal::new(CrosstermBackend::new()).unwrap();
    terminal.clear().unwrap();
    terminal.hide_cursor().unwrap();
    draw(&mut terminal).unwrap();
    loop {
        {
            let input = crossterm::input(terminal.backend().screen());
            match input.read_char() {
                Ok(c) => if c == 'q' {
                    break;
                },
                Err(e) => panic!("{}", e.description()),
            };
        }
        draw(&mut terminal).unwrap();
    }
    terminal.show_cursor().unwrap();
}

fn draw(t: &mut Terminal<CrosstermBackend>) -> io::Result<()> {
    let size = t.size()?;
    t.draw(|mut f| {
        let text = [
            Text::Data("It "),
            Text::StyledData("works", Style::default().fg(Color::Yellow)),
        ];
        Paragraph::new(text.iter())
            .block(
                Block::default()
                    .title("Crossterm Backend")
                    .title_style(Style::default().fg(Color::Yellow).modifier(Modifier::Bold))
                    .borders(Borders::ALL)
                    .border_style(Style::default().fg(Color::Magenta)),
            )
            .render(&mut f, size);
    })
}
<commit_msg>fix: update crossterm example<commit_after>extern crate crossterm;
extern crate tui;

use std::error::Error;
use std::io;

use tui::backend::CrosstermBackend;
use tui::layout::{Constraint, Direction, Layout};
use tui::style::{Color, Modifier, Style};
use tui::widgets::{Block, Borders, Paragraph, Text, Widget};
use tui::Terminal;

fn main() {
    let mut terminal = Terminal::new(CrosstermBackend::new()).unwrap();
    terminal.clear().unwrap();
    terminal.hide_cursor().unwrap();
    draw(&mut terminal).unwrap();
    loop {
        {
            let input = crossterm::input(terminal.backend().screen());
            match input.read_char() {
                Ok(c) => if c == 'q' {
                    break;
                },
                Err(e) => panic!("{}", e.description()),
            };
        }
        draw(&mut terminal).unwrap();
    }
    terminal.show_cursor().unwrap();
}

fn draw(t: &mut Terminal<CrosstermBackend>) -> io::Result<()> {
    let size = t.size()?;
    t.draw(|mut f| {
        let text = [
            Text::raw("It "),
            Text::styled("works", Style::default().fg(Color::Yellow)),
        ];
        Paragraph::new(text.iter())
            .block(
                Block::default()
                    .title("Crossterm Backend")
                    .title_style(Style::default().fg(Color::Yellow).modifier(Modifier::Bold))
                    .borders(Borders::ALL)
                    .border_style(Style::default().fg(Color::Magenta)),
            )
            .render(&mut f, size);
    })
}
<|endoftext|>


<commit_before>// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{borrow::Cow, collections::HashSet, fmt::Display, io, path::PathBuf, process};

use anyhow::Error;
use clap::{crate_authors, crate_version, Arg, ArgMatches, Command};
use itertools::Itertools;
use mdbook::{book::Book, preprocess::CmdPreprocessor};
use proc_macro2::Span;
use syn::{spanned::Spanned, Expr, __private::ToTokens};

static LONG_ABOUT: &str =
    "This is an mdbook preprocessor tailored for autocxx code examples. Autocxx
code examples don't fit well 'mdbook test' or even alternatives such as
'skeptic' or 'doc_comment' for these reasons:

a) A single code example consists of both Rust and C++ code. They must be
   linked into a separate executable, i.e. we must make one executable per
   doc test.
b) The code examples must be presented/formatted nicely with suitable
   separate blocks for the Rust and C++ code.
c) mdbook test is not good at handling doctests which have dependencies.

This preprocessor will find code snippets like this:
```rust,autocxx
autocxx_integration_tests::doctest(
\" /* any C++ implementation code */\",
\" /* C++ header code */\",
{
/* complete Rust code including 'main' */
)
```

and will build and run them, while emitting better formatted markdown blocks
for subsequent preprocessors and renderers.
";

fn main() {
    let matches = Command::new("autocxx-mdbook-preprocessor")
        .version(crate_version!())
        .author(crate_authors!())
        .about("Expands and tests code examples in the autocxx book.")
        .long_about(LONG_ABOUT)
        .subcommand(
            Command::new("supports")
                .arg(Arg::new("renderer").required(true))
                .about("Whether a given renderer is supported by this preprocessor"),
        )
        .arg(
            Arg::new("skip_tests")
                .short('s')
                .help("Skip running doctests"),
        )
        .arg(
            Arg::new("manifest_dir")
            .long("manifest-dir")
            .help("Path to directory containing outermost autocxx Cargo.toml; necessary for trybuild to build test code successfully")
            .default_value_os(calculate_cargo_dir().as_os_str())
        )
        .get_matches();
    if matches.subcommand_matches("supports").is_some() {
        process::exit(0); // we support all renderers.
    }
    preprocess(&matches).unwrap();
}

fn calculate_cargo_dir() -> PathBuf {
    let mut path = std::env::current_exe().unwrap();
    for _ in 0..3 {
        path = path.parent().map(|p| p.to_path_buf()).unwrap_or(path);
    }
    path.join("integration-tests")
}

fn preprocess(args: &ArgMatches) -> Result<(), Error> {
    let (_, mut book) = CmdPreprocessor::parse_input(io::stdin())?;

    let mut test_cases = Vec::new();

    Book::for_each_mut(&mut book, |sec| {
        if let mdbook::BookItem::Chapter(chapter) = sec {
            let filename = chapter
                .path
                .as_ref()
                .map(|pb| pb.to_string_lossy())
                .unwrap_or_default()
                .to_string();
            chapter.content = substitute_chapter(&chapter.content, &filename, &mut test_cases);
        }
    });

    // Now run any test cases we accumulated.
    if !args.is_present("skip_tests") {
        let num_tests = test_cases.len();
        for (counter, case) in test_cases.into_iter().enumerate() {
            eprintln!(
                "Running doctest {}/{} at {}",
                counter + 1,
                num_tests,
                &case.location
            );
            autocxx_integration_tests::doctest(
                &case.cpp,
                &case.hdr,
                case.rs.to_token_stream(),
                args.value_of_os("manifest_dir").unwrap(),
            );
            eprintln!(
                "Doctest {}/{} at {} passed.",
                counter + 1,
                num_tests,
                &case.location
            );
        }
    }

    serde_json::to_writer(io::stdout(), &book)?;

    Ok(())
}

fn substitute_chapter(chapter: &str, filename: &str, test_cases: &mut Vec<TestCase>) -> String {
    let mut state = ChapterParseState::Start;
    let mut out = Vec::new();
    for (line_no, line) in chapter.lines().enumerate() {
        let line_type = recognize_line(line);
        let mut push_line = true;
        state = match state {
            ChapterParseState::Start => match line_type {
                LineType::CodeBlockStart | LineType::CodeBlockEnd => {
                    ChapterParseState::OtherCodeBlock
                }
                LineType::CodeBlockStartAutocxx(block_flags) => {
                    push_line = false;
                    ChapterParseState::OurCodeBlock(block_flags, Vec::new())
                }
                LineType::Misc => ChapterParseState::Start,
            },
            ChapterParseState::OtherCodeBlock => match line_type {
                LineType::CodeBlockEnd => ChapterParseState::Start,
                LineType::Misc => ChapterParseState::OtherCodeBlock,
                _ => panic!("Found confusing conflicting block markers"),
            },
            ChapterParseState::OurCodeBlock(flags, mut lines) => match line_type {
                LineType::Misc => {
                    push_line = false;
                    lines.push(line.to_string());
                    ChapterParseState::OurCodeBlock(flags, lines)
                }
                LineType::CodeBlockEnd => {
                    let location = MiniSpan {
                        filename: filename.to_string(),
                        start_line: line_no - lines.len(),
                    };
                    out.extend(handle_code_block(flags, lines, location, test_cases));
                    push_line = false;
                    ChapterParseState::Start
                }
                _ => panic!("Found something unexpected in one of our code blocks"),
            },
        };
        if push_line {
            out.push(line.to_string());
        }
    }

    out.join("\n")
}

/// Like `proc_macro2::Span` but only has the starting line. For basic
/// diagnostics.
struct MiniSpan {
    filename: String,
    start_line: usize,
}

impl Display for MiniSpan {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{} line {}", self.filename, self.start_line)
    }
}

struct TestCase {
    cpp: String,
    hdr: String,
    rs: syn::File,
    location: MiniSpan,
}

enum ChapterParseState {
    Start,
    OtherCodeBlock,
    OurCodeBlock(HashSet<String>, Vec<String>), // have found rust,autocxx
}

enum LineType {
    CodeBlockStart,
    CodeBlockStartAutocxx(HashSet<String>),
    CodeBlockEnd,
    Misc,
}

fn code_block_flags(line: &str) -> HashSet<String> {
    let line = &line[3..];
    line.split(',').map(|s| s.to_string()).collect()
}

fn recognize_line(line: &str) -> LineType {
    if line.starts_with("```") && line.len() > 3 {
        let flags = code_block_flags(line);
        if flags.contains("autocxx") {
            LineType::CodeBlockStartAutocxx(flags)
        } else {
            LineType::CodeBlockStart
        }
    } else if line == "```" {
        LineType::CodeBlockEnd
    } else {
        LineType::Misc
    }
}

fn handle_code_block(
    flags: HashSet<String>,
    lines: Vec<String>,
    location: MiniSpan,
    test_cases: &mut Vec<TestCase>,
) -> impl Iterator<Item = String> {
    let input_str = lines.join("\n");
    let fn_call = syn::parse_str::<syn::Expr>(&input_str).expect("Unable to parse");
    let fn_call = match fn_call {
        Expr::Call(expr) => expr,
        _ => panic!("Parsing unexpected"),
    };
    let mut args_iter = fn_call.args.iter();
    let cpp = extract_span(&lines, args_iter.next().unwrap().span());
    let hdr = extract_span(&lines, args_iter.next().unwrap().span());
    let rs = extract_span(&lines, args_iter.next().unwrap().span());
    let mut output = vec![
        "#### C++ header:".to_string(),
        "```cpp".to_string(),
        hdr.to_string(),
        "```".to_string(),
    ];
    if !cpp.is_empty() && !flags.contains("hidecpp") {
        output.push("#### C++ implementation:".to_string());
        output.push("```cpp".to_string());
        output.push(cpp.to_string());
        output.push("```".to_string());
    }
    output.push("#### Rust:".to_string());
    output.push("```rust,noplayground".to_string());
    output.push(escape_hexathorpes(&rs).to_string());
    output.push("```".to_string());

    // Don't run the test cases yet, because we want the preprocessor to spot
    // basic formatting errors before getting into the time consuming business of
    // running tests.
    if !flags.contains("nocompile") {
        test_cases.push(TestCase {
            cpp: cpp.to_string(),
            hdr: hdr.to_string(),
            rs: syn::parse_file(&rs).unwrap(),
            location,
        });
    }

    output.into_iter()
}

fn extract_span(text: &[String], span: Span) -> Cow<str> {
    let start_line = span.start().line - 1;
    let start_col = span.start().column;
    let end_line = span.end().line - 1;
    let end_col = span.end().column;
    if start_line == end_line {
        Cow::Borrowed(&text[start_line][start_col + 1..end_col - 1])
    } else {
        let start_subset = &text[start_line][start_col + 1..];
        let end_subset = &text[end_line][..end_col - 1];
        let mid_lines = &text[start_line + 1..end_line];
        Cow::Owned(
            std::iter::once(start_subset.to_string())
                .chain(mid_lines.iter().cloned())
                .chain(std::iter::once(end_subset.to_string()))
                .join("\n"),
        )
    }
}

fn escape_hexathorpes(input: &str) -> Cow<str> {
    let re = regex::Regex::new(r"(?m)^(?P<ws>\s*)#(?P<c>.*)").unwrap();
    re.replace_all(input, "$ws##$c")
}
<commit_msg>Unescape quotes into C++<commit_after>// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{borrow::Cow, collections::HashSet, fmt::Display, io, path::PathBuf, process};

use anyhow::Error;
use clap::{crate_authors, crate_version, Arg, ArgMatches, Command};
use itertools::Itertools;
use mdbook::{book::Book, preprocess::CmdPreprocessor};
use proc_macro2::Span;
use syn::{spanned::Spanned, Expr, __private::ToTokens};

static LONG_ABOUT: &str =
    "This is an mdbook preprocessor tailored for autocxx code examples. Autocxx
code examples don't fit well 'mdbook test' or even alternatives such as
'skeptic' or 'doc_comment' for these reasons:

a) A single code example consists of both Rust and C++ code. They must be
   linked into a separate executable, i.e. we must make one executable per
   doc test.
b) The code examples must be presented/formatted nicely with suitable
   separate blocks for the Rust and C++ code.
c) mdbook test is not good at handling doctests which have dependencies.

This preprocessor will find code snippets like this:
```rust,autocxx
autocxx_integration_tests::doctest(
\" /* any C++ implementation code */\",
\" /* C++ header code */\",
{
/* complete Rust code including 'main' */
)
```

and will build and run them, while emitting better formatted markdown blocks
for subsequent preprocessors and renderers.
";

fn main() {
    let matches = Command::new("autocxx-mdbook-preprocessor")
        .version(crate_version!())
        .author(crate_authors!())
        .about("Expands and tests code examples in the autocxx book.")
        .long_about(LONG_ABOUT)
        .subcommand(
            Command::new("supports")
                .arg(Arg::new("renderer").required(true))
                .about("Whether a given renderer is supported by this preprocessor"),
        )
        .arg(
            Arg::new("skip_tests")
                .short('s')
                .help("Skip running doctests"),
        )
        .arg(
            Arg::new("manifest_dir")
            .long("manifest-dir")
            .help("Path to directory containing outermost autocxx Cargo.toml; necessary for trybuild to build test code successfully")
            .default_value_os(calculate_cargo_dir().as_os_str())
        )
        .get_matches();
    if matches.subcommand_matches("supports").is_some() {
        process::exit(0); // we support all renderers.
    }
    preprocess(&matches).unwrap();
}

fn calculate_cargo_dir() -> PathBuf {
    let mut path = std::env::current_exe().unwrap();
    for _ in 0..3 {
        path = path.parent().map(|p| p.to_path_buf()).unwrap_or(path);
    }
    path.join("integration-tests")
}

fn preprocess(args: &ArgMatches) -> Result<(), Error> {
    let (_, mut book) = CmdPreprocessor::parse_input(io::stdin())?;

    let mut test_cases = Vec::new();

    Book::for_each_mut(&mut book, |sec| {
        if let mdbook::BookItem::Chapter(chapter) = sec {
            let filename = chapter
                .path
                .as_ref()
                .map(|pb| pb.to_string_lossy())
                .unwrap_or_default()
                .to_string();
            chapter.content = substitute_chapter(&chapter.content, &filename, &mut test_cases);
        }
    });

    // Now run any test cases we accumulated.
    if !args.is_present("skip_tests") {
        let num_tests = test_cases.len();
        for (counter, case) in test_cases.into_iter().enumerate() {
            eprintln!(
                "Running doctest {}/{} at {}",
                counter + 1,
                num_tests,
                &case.location
            );
            autocxx_integration_tests::doctest(
                &case.cpp,
                &case.hdr,
                case.rs.to_token_stream(),
                args.value_of_os("manifest_dir").unwrap(),
            );
            eprintln!(
                "Doctest {}/{} at {} passed.",
                counter + 1,
                num_tests,
                &case.location
            );
        }
    }

    serde_json::to_writer(io::stdout(), &book)?;

    Ok(())
}

fn substitute_chapter(chapter: &str, filename: &str, test_cases: &mut Vec<TestCase>) -> String {
    let mut state = ChapterParseState::Start;
    let mut out = Vec::new();
    for (line_no, line) in chapter.lines().enumerate() {
        let line_type = recognize_line(line);
        let mut push_line = true;
        state = match state {
            ChapterParseState::Start => match line_type {
                LineType::CodeBlockStart | LineType::CodeBlockEnd => {
                    ChapterParseState::OtherCodeBlock
                }
                LineType::CodeBlockStartAutocxx(block_flags) => {
                    push_line = false;
                    ChapterParseState::OurCodeBlock(block_flags, Vec::new())
                }
                LineType::Misc => ChapterParseState::Start,
            },
            ChapterParseState::OtherCodeBlock => match line_type {
                LineType::CodeBlockEnd => ChapterParseState::Start,
                LineType::Misc => ChapterParseState::OtherCodeBlock,
                _ => panic!("Found confusing conflicting block markers"),
            },
            ChapterParseState::OurCodeBlock(flags, mut lines) => match line_type {
                LineType::Misc => {
                    push_line = false;
                    lines.push(line.to_string());
                    ChapterParseState::OurCodeBlock(flags, lines)
                }
                LineType::CodeBlockEnd => {
                    let location = MiniSpan {
                        filename: filename.to_string(),
                        start_line: line_no - lines.len(),
                    };
                    out.extend(handle_code_block(flags, lines, location, test_cases));
                    push_line = false;
                    ChapterParseState::Start
                }
                _ => panic!("Found something unexpected in one of our code blocks"),
            },
        };
        if push_line {
            out.push(line.to_string());
        }
    }

    out.join("\n")
}

/// Like `proc_macro2::Span` but only has the starting line. For basic
/// diagnostics.
struct MiniSpan {
    filename: String,
    start_line: usize,
}

impl Display for MiniSpan {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{} line {}", self.filename, self.start_line)
    }
}

struct TestCase {
    cpp: String,
    hdr: String,
    rs: syn::File,
    location: MiniSpan,
}

enum ChapterParseState {
    Start,
    OtherCodeBlock,
    OurCodeBlock(HashSet<String>, Vec<String>), // have found rust,autocxx
}

enum LineType {
    CodeBlockStart,
    CodeBlockStartAutocxx(HashSet<String>),
    CodeBlockEnd,
    Misc,
}

fn code_block_flags(line: &str) -> HashSet<String> {
    let line = &line[3..];
    line.split(',').map(|s| s.to_string()).collect()
}

fn recognize_line(line: &str) -> LineType {
    if line.starts_with("```") && line.len() > 3 {
        let flags = code_block_flags(line);
        if flags.contains("autocxx") {
            LineType::CodeBlockStartAutocxx(flags)
        } else {
            LineType::CodeBlockStart
        }
    } else if line == "```" {
        LineType::CodeBlockEnd
    } else {
        LineType::Misc
    }
}

fn handle_code_block(
    flags: HashSet<String>,
    lines: Vec<String>,
    location: MiniSpan,
    test_cases: &mut Vec<TestCase>,
) -> impl Iterator<Item = String> {
    let input_str = lines.join("\n");
    let fn_call = syn::parse_str::<syn::Expr>(&input_str).expect("Unable to parse");
    let fn_call = match fn_call {
        Expr::Call(expr) => expr,
        _ => panic!("Parsing unexpected"),
    };
    let mut args_iter = fn_call.args.iter();
    let cpp = extract_span(&lines, args_iter.next().unwrap().span());
    let hdr = extract_span(&lines, args_iter.next().unwrap().span());
    let rs = extract_span(&lines, args_iter.next().unwrap().span());
    let mut output = vec![
        "#### C++ header:".to_string(),
        "```cpp".to_string(),
        hdr.to_string(),
        "```".to_string(),
    ];
    if !cpp.is_empty() && !flags.contains("hidecpp") {
        output.push("#### C++ implementation:".to_string());
        output.push("```cpp".to_string());
        output.push(cpp.to_string());
        output.push("```".to_string());
    }
    output.push("#### Rust:".to_string());
    output.push("```rust,noplayground".to_string());
    output.push(escape_hexathorpes(&rs).to_string());
    output.push("```".to_string());

    // Don't run the test cases yet, because we want the preprocessor to spot
    // basic formatting errors before getting into the time consuming business of
    // running tests.
    if !flags.contains("nocompile") {
        test_cases.push(TestCase {
            cpp: unescape_quotes(&cpp),
            hdr: unescape_quotes(&hdr),
            rs: syn::parse_file(&rs).unwrap(),
            location,
        });
    }

    output.into_iter()
}

fn extract_span(text: &[String], span: Span) -> Cow<str> {
    let start_line = span.start().line - 1;
    let start_col = span.start().column;
    let end_line = span.end().line - 1;
    let end_col = span.end().column;
    if start_line == end_line {
        Cow::Borrowed(&text[start_line][start_col + 1..end_col - 1])
    } else {
        let start_subset = &text[start_line][start_col + 1..];
        let end_subset = &text[end_line][..end_col - 1];
        let mid_lines = &text[start_line + 1..end_line];
        Cow::Owned(
            std::iter::once(start_subset.to_string())
                .chain(mid_lines.iter().cloned())
                .chain(std::iter::once(end_subset.to_string()))
                .join("\n"),
        )
    }
}

fn escape_hexathorpes(input: &str) -> Cow<str> {
    let re = regex::Regex::new(r"(?m)^(?P<ws>\s*)#(?P<c>.*)").unwrap();
    re.replace_all(input, "$ws##$c")
}

fn unescape_quotes(input: &str) -> String {
    input.replace("\\\"", "\"")
}
<|endoftext|>