<commit_before>use nom::{
  Consumer, ConsumerState,
  ErrorKind,
  HexDisplay,
  Input, IResult,
  Move, Needed,
};
use metadata::parser::{header, block_data};

use std::collections::HashMap;

/// Data associated with a single metadata block.
#[derive(Debug)]
pub struct Metadata {
  /// Marks whether the current metadata block is the last.
  pub is_last: bool,
  /// The length, in bytes, of the block being parsed. This does not include
  /// the metadata block header.
  pub length: u32,
  /// Block data containing one of the eight different types of metadata.
  pub data: Data,
}

/// General enum that hold all the different metadata block data.
#[derive(Debug, PartialEq, Eq)]
pub enum Data {
  /// Information regarding the entire audio stream.
  StreamInfo(StreamInfo),
  /// Block that represents a number of padded bytes.
  Padding(u32),
  /// Data used by third-party applications.
  Application(Application),
  /// Table of multiple points to seek, or skip, to within the FLAC file.
  SeekTable(Vec<SeekPoint>),
  /// Stores human-readable name/value pairs.
  VorbisComment(VorbisComment),
  /// Stores cue information
  CueSheet(CueSheet),
  /// Stores pictures associated with the FLAC file.
  Picture(Picture),
  /// A type of block data that isn't know or doesn't match the type above.
  Unknown(Vec<u8>),
}

/// Information regarding the entire audio stream.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct StreamInfo {
  /// Minimum block size, in samples, used in the stream.
  pub min_block_size: u16,
  /// Maximum block size, in samples, used in the stream.
  pub max_block_size: u16,
  /// Minimum frame size, in bytes, used in the stream. May be zero to imply
  /// the value isn't know.
  pub min_frame_size: u32,
  /// Maximum frame size, in bytes, used in the stream. May be zero to imply
  /// the value isn't know.
  pub max_frame_size: u32,
  /// Sample rate in hertz (Hz).
  pub sample_rate: u32,
  /// Number of channels. FLAC supports one to eight channel.
  pub channels: u8,
  /// Bits per sample. FLAC supports four to thirty-two bits per sample.
  pub bits_per_sample: u8,
  /// Total samples in the stream. A value of zero means the number is
  /// unknown.
  pub total_samples: u64,
  /// MD5 signature of the unencoded audio data.
  pub md5_sum: [u8; 16],
}

impl StreamInfo {
  pub fn new() -> StreamInfo {
    StreamInfo {
      min_block_size: 0,
      max_block_size: 0,
      min_frame_size: 0,
      max_frame_size: 0,
      sample_rate: 0,
      channels: 0,
      bits_per_sample: 0,
      total_samples: 0,
      md5_sum: [0; 16],
    }
  }
}

/// Data used by third-party applications.
#[derive(Debug, PartialEq, Eq)]
pub struct Application {
  /// Registered application ID.
  pub id: String,
  /// Data used by the third-party application.
  pub data: Vec<u8>,
}

/// Seek, or skip, to a point within the FLAC file.
#[derive(Debug, PartialEq, Eq)]
pub struct SeekPoint {
  /// Sample number of the first sample in the target frame.
  pub sample_number: u64,
  /// Byte offset of the target frame's header.
  pub stream_offset: u64,
  /// Number of samples in the target frame.
  pub frame_samples: u16,
}

/// Stores human-readable name/value pairs.
#[derive(Debug, PartialEq, Eq)]
pub struct VorbisComment {
  /// Vendor name.
  pub vendor_string: String,
  /// Comments associated with a name, or category, followed by it's
  /// contents.
  pub comments: HashMap<String, String>,
}

/// Stores cue information.
///
/// Generally for storing information from Compact Disk Digital Audio, but
/// can be used as a cueing mechanism for playback.
#[derive(Debug, PartialEq, Eq)]
pub struct CueSheet {
  /// Media catalog number.
  pub media_catalog_number: String,
  /// Number of lead-in samples.
  pub lead_in: u64,
  /// Whether or not this `CueSheet` corresponds to a Compact Disc.
  pub is_cd: bool,
  /// One or more tracks.
  pub tracks: Vec<CueSheetTrack>,
}

/// Track information inside a cue sheet.
#[derive(Debug, PartialEq, Eq)]
pub struct CueSheetTrack {
  /// Track offset, in samples, relative to the beginning of the FLAC audio
  /// stream.
  pub offset: u64,
  /// Track number.
  pub number: u8,
  /// Twelve digit alphanumeric code.
  pub isrc: String,
  /// Whether the cue sheet track is audio.
  pub is_audio: bool,
  /// Whether the cue sheet track is pre-emphasis.
  pub is_pre_emphasis: bool,
  /// For all tracks except the lead-out track.
  pub indices: Vec<CueSheetTrackIndex>,
}

/// An index point within a track, inside of a cue sheet.
#[derive(Debug, PartialEq, Eq)]
pub struct CueSheetTrackIndex {
  /// Offset, in samples, relative to the track offset of the index point.
  pub offset: u64,
  /// Index point number.
  pub number: u8,
}

/// Stores pictures associated with the FLAC file.
///
/// More than likely these pictures will be cover art, but you can have more
/// than one within a file, which are distinguished by `PictureType`and it's
/// mime type string.
#[derive(Debug, PartialEq, Eq)]
pub struct Picture {
  /// Picture type, based on the ID3v2 APIC frame.
  pub picture_type: PictureType,
  /// Multipurpose Internet Mail Extensions (MIME) type.
  pub mime_type: String,
  /// A string describing the picture.
  pub description: String,
  /// Width of the picture in pixels.
  pub width: u32,
  /// Height of the picture in pixels.
  pub height: u32,
  /// Color depth of the picture in bits-per-pixel.
  pub depth: u32,
  /// Number of colors used.
  pub colors: u32,
  /// Binary picture data.
  pub data: Vec<u8>,
}

/// The picture type according to the ID3v2 attached picture frame.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum PictureType {
  /// Other picture type not categorized in this enum.
  Other,
  /// 32x32 pixels 'file icon'.
  FileIconStandard,
  /// Other, or non-standard, file icon.
  FileIcon,
  /// Cover (front).
  FrontCover,
  /// Cover (back).
  BackCover,
  /// Leaflet page.
  LeafletPage,
  /// Media, like label side of a CD.
  Media,
  /// Lead artist, lead performer, or soloist.
  LeadArtist,
  /// Artist or performer.
  Artist,
  /// Conductor.
  Conductor,
  /// Band or orchestra.
  Band,
  /// Composer.
  Composer,
  /// Lyricist or text writer.
  Lyricist,
  /// Recording location.
  RecordingLocation,
  /// During recording.
  DuringRecording,
  /// During performance.
  DuringPerformace,
  /// Movie, or video, screen capture.
  VideoScreenCapture,
  /// A bright colored fish.
  Fish,
  /// Illustration.
  Illustration,
  /// Band, or artist, logotype.
  BandLogo,
  /// Publisher, or studio, logotype.
  PublisherLogo,
}

enum ParserState {
  Marker,
  Header,
  Block((bool, u8, u32)),
}

pub struct MetaDataConsumer {
  state: ParserState,
  consumer_state: ConsumerState<(), ErrorKind, Move>,
  pub data: Vec<Metadata>,
}

impl MetaDataConsumer {
  pub fn new() -> MetaDataConsumer {
    let consumed = Move::Consume(0);

    MetaDataConsumer {
      state: ParserState::Marker,
      consumer_state: ConsumerState::Continue(consumed),
      data: Vec::new(),
    }
  }

  fn handle_marker<'a>(&mut self, input: &'a [u8]) -> IResult<&'a [u8], ()> {
    let kind = ErrorKind::Custom(0);

    match tag!(input, "fLaC") {
      IResult::Done(i, _)    => {
        self.state = ParserState::Header;

        IResult::Error(Err::Position(kind, i))
      }
      IResult::Error(_)      => IResult::Error(Err::Code(kind)),
      IResult::Incomplete(n) => IResult::Incomplete(n),
    }
  }

  fn handle_header<'a>(&mut self, input: &'a [u8]) -> IResult<&'a [u8], ()> {
    match header(input) {
      IResult::Done(i, data) => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);

        self.state          = ParserState::Block(data);
        self.consumer_state = ConsumerState::Continue(consumed);
      }
      IResult::Error(_)      => {
        let kind = ErrorKind::Custom(1);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(_) => {
        let needed = Move::Await(Needed::Size(4));

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }

  fn handle_block<'a>(&mut self, input: &'a [u8], header: (bool, u8, u32))
                      -> IResult<&'a [u8], ()> {
    let (is_last, block_type, length) = header;

    match block_data(input, block_type, length) {
      IResult::Done(i, data) => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);

        self.data.push(Metadata {
          is_last: is_last,
          length: length,
          data: data,
        });

        if is_last {
          self.consumer_state = ConsumerState::Done(consumed, ());
        } else {
          self.state          = ParserState::Header;
          self.consumer_state = ConsumerState::Continue(consumed);
        }
      }
      IResult::Error(_)      => {
        let kind = ErrorKind::Custom(2);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(_) => {
        let needed = Move::Await(Needed::Size(length as usize));

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }
}

impl<'a> Consumer<&'a [u8], (), ErrorKind, Move> for MetaDataConsumer {
  fn state(&self) -> &ConsumerState<(), ErrorKind, Move> {
    &self.consumer_state
  }

  fn handle(&mut self, input: Input<&'a [u8]>)
            -> &ConsumerState<(), ErrorKind, Move> {
    match input {
      Input::Element(i) | Input::Eof(Some(i)) => {
        match self.state {
          ParserState::Marker      => self.handle_marker(i),
          ParserState::Header      => self.handle_header(i),
          ParserState::Block(data) => self.handle_block(i, data),
        }
      }
      Input::Empty | Input::Eof(None)         => {
        let kind = match self.state {
          ParserState::Marker   => ErrorKind::Custom(0),
          ParserState::Header   => ErrorKind::Custom(1),
          ParserState::Block(_) => ErrorKind::Custom(2),
        };

        self.consumer_state = ConsumerState::Error(kind);
      }
    }

    &self.consumer_state
  }
}
<commit_msg>Rewrite `handle_header` for what it actually need<commit_after>use nom::{
  Consumer, ConsumerState,
  ErrorKind,
  HexDisplay,
  Input, IResult,
  Move, Needed,
};
use metadata::parser::{header, block_data};

use std::collections::HashMap;

/// Data associated with a single metadata block.
#[derive(Debug)]
pub struct Metadata {
  /// Marks whether the current metadata block is the last.
  pub is_last: bool,
  /// The length, in bytes, of the block being parsed. This does not include
  /// the metadata block header.
  pub length: u32,
  /// Block data containing one of the eight different types of metadata.
  pub data: Data,
}

/// General enum that hold all the different metadata block data.
#[derive(Debug, PartialEq, Eq)]
pub enum Data {
  /// Information regarding the entire audio stream.
  StreamInfo(StreamInfo),
  /// Block that represents a number of padded bytes.
  Padding(u32),
  /// Data used by third-party applications.
  Application(Application),
  /// Table of multiple points to seek, or skip, to within the FLAC file.
  SeekTable(Vec<SeekPoint>),
  /// Stores human-readable name/value pairs.
  VorbisComment(VorbisComment),
  /// Stores cue information
  CueSheet(CueSheet),
  /// Stores pictures associated with the FLAC file.
  Picture(Picture),
  /// A type of block data that isn't know or doesn't match the type above.
  Unknown(Vec<u8>),
}

/// Information regarding the entire audio stream.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct StreamInfo {
  /// Minimum block size, in samples, used in the stream.
  pub min_block_size: u16,
  /// Maximum block size, in samples, used in the stream.
  pub max_block_size: u16,
  /// Minimum frame size, in bytes, used in the stream. May be zero to imply
  /// the value isn't know.
  pub min_frame_size: u32,
  /// Maximum frame size, in bytes, used in the stream. May be zero to imply
  /// the value isn't know.
  pub max_frame_size: u32,
  /// Sample rate in hertz (Hz).
  pub sample_rate: u32,
  /// Number of channels. FLAC supports one to eight channel.
  pub channels: u8,
  /// Bits per sample. FLAC supports four to thirty-two bits per sample.
  pub bits_per_sample: u8,
  /// Total samples in the stream. A value of zero means the number is
  /// unknown.
  pub total_samples: u64,
  /// MD5 signature of the unencoded audio data.
  pub md5_sum: [u8; 16],
}

impl StreamInfo {
  pub fn new() -> StreamInfo {
    StreamInfo {
      min_block_size: 0,
      max_block_size: 0,
      min_frame_size: 0,
      max_frame_size: 0,
      sample_rate: 0,
      channels: 0,
      bits_per_sample: 0,
      total_samples: 0,
      md5_sum: [0; 16],
    }
  }
}

/// Data used by third-party applications.
#[derive(Debug, PartialEq, Eq)]
pub struct Application {
  /// Registered application ID.
  pub id: String,
  /// Data used by the third-party application.
  pub data: Vec<u8>,
}

/// Seek, or skip, to a point within the FLAC file.
#[derive(Debug, PartialEq, Eq)]
pub struct SeekPoint {
  /// Sample number of the first sample in the target frame.
  pub sample_number: u64,
  /// Byte offset of the target frame's header.
  pub stream_offset: u64,
  /// Number of samples in the target frame.
  pub frame_samples: u16,
}

/// Stores human-readable name/value pairs.
#[derive(Debug, PartialEq, Eq)]
pub struct VorbisComment {
  /// Vendor name.
  pub vendor_string: String,
  /// Comments associated with a name, or category, followed by it's
  /// contents.
  pub comments: HashMap<String, String>,
}

/// Stores cue information.
///
/// Generally for storing information from Compact Disk Digital Audio, but
/// can be used as a cueing mechanism for playback.
#[derive(Debug, PartialEq, Eq)]
pub struct CueSheet {
  /// Media catalog number.
  pub media_catalog_number: String,
  /// Number of lead-in samples.
  pub lead_in: u64,
  /// Whether or not this `CueSheet` corresponds to a Compact Disc.
  pub is_cd: bool,
  /// One or more tracks.
  pub tracks: Vec<CueSheetTrack>,
}

/// Track information inside a cue sheet.
#[derive(Debug, PartialEq, Eq)]
pub struct CueSheetTrack {
  /// Track offset, in samples, relative to the beginning of the FLAC audio
  /// stream.
  pub offset: u64,
  /// Track number.
  pub number: u8,
  /// Twelve digit alphanumeric code.
  pub isrc: String,
  /// Whether the cue sheet track is audio.
  pub is_audio: bool,
  /// Whether the cue sheet track is pre-emphasis.
  pub is_pre_emphasis: bool,
  /// For all tracks except the lead-out track.
  pub indices: Vec<CueSheetTrackIndex>,
}

/// An index point within a track, inside of a cue sheet.
#[derive(Debug, PartialEq, Eq)]
pub struct CueSheetTrackIndex {
  /// Offset, in samples, relative to the track offset of the index point.
  pub offset: u64,
  /// Index point number.
  pub number: u8,
}

/// Stores pictures associated with the FLAC file.
///
/// More than likely these pictures will be cover art, but you can have more
/// than one within a file, which are distinguished by `PictureType`and it's
/// mime type string.
#[derive(Debug, PartialEq, Eq)]
pub struct Picture {
  /// Picture type, based on the ID3v2 APIC frame.
  pub picture_type: PictureType,
  /// Multipurpose Internet Mail Extensions (MIME) type.
  pub mime_type: String,
  /// A string describing the picture.
  pub description: String,
  /// Width of the picture in pixels.
  pub width: u32,
  /// Height of the picture in pixels.
  pub height: u32,
  /// Color depth of the picture in bits-per-pixel.
  pub depth: u32,
  /// Number of colors used.
  pub colors: u32,
  /// Binary picture data.
  pub data: Vec<u8>,
}

/// The picture type according to the ID3v2 attached picture frame.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum PictureType {
  /// Other picture type not categorized in this enum.
  Other,
  /// 32x32 pixels 'file icon'.
  FileIconStandard,
  /// Other, or non-standard, file icon.
  FileIcon,
  /// Cover (front).
  FrontCover,
  /// Cover (back).
  BackCover,
  /// Leaflet page.
  LeafletPage,
  /// Media, like label side of a CD.
  Media,
  /// Lead artist, lead performer, or soloist.
  LeadArtist,
  /// Artist or performer.
  Artist,
  /// Conductor.
  Conductor,
  /// Band or orchestra.
  Band,
  /// Composer.
  Composer,
  /// Lyricist or text writer.
  Lyricist,
  /// Recording location.
  RecordingLocation,
  /// During recording.
  DuringRecording,
  /// During performance.
  DuringPerformace,
  /// Movie, or video, screen capture.
  VideoScreenCapture,
  /// A bright colored fish.
  Fish,
  /// Illustration.
  Illustration,
  /// Band, or artist, logotype.
  BandLogo,
  /// Publisher, or studio, logotype.
  PublisherLogo,
}

enum ParserState {
  Marker,
  Header,
  Block((bool, u8, u32)),
}

pub struct MetaDataConsumer {
  state: ParserState,
  consumer_state: ConsumerState<(), ErrorKind, Move>,
  pub data: Vec<Metadata>,
}

impl MetaDataConsumer {
  pub fn new() -> MetaDataConsumer {
    let consumed = Move::Consume(0);

    MetaDataConsumer {
      state: ParserState::Marker,
      consumer_state: ConsumerState::Continue(consumed),
      data: Vec::new(),
    }
  }

  fn handle_marker<'a>(&mut self, input: &'a [u8]) -> IResult<&'a [u8], ()> {
    let kind = ErrorKind::Custom(0);

    match tag!(input, "fLaC") {
      IResult::Done(i, _)    => {
        self.state = ParserState::Header;

        IResult::Error(Err::Position(kind, i))
      }
      IResult::Error(_)      => IResult::Error(Err::Code(kind)),
      IResult::Incomplete(n) => IResult::Incomplete(n),
    }
  }

  fn handle_header<'a>(&mut self, input: &'a [u8]) -> IResult<&'a [u8], ()> {
    let kind = ErrorKind::Custom(1);

    match header(input) {
      IResult::Done(i, data) => {
        self.state = ParserState::Block(data);

        IResult::Error(Err::Position(kind, i))
      }
      IResult::Error(_)      => IResult::Error(Err::Code(kind)),
      IResult::Incomplete(n) => IResult::Incomplete(n),
    }
  }

  fn handle_block<'a>(&mut self, input: &'a [u8], header: (bool, u8, u32))
                      -> IResult<&'a [u8], ()> {
    let (is_last, block_type, length) = header;

    match block_data(input, block_type, length) {
      IResult::Done(i, data) => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);

        self.data.push(Metadata {
          is_last: is_last,
          length: length,
          data: data,
        });

        if is_last {
          self.consumer_state = ConsumerState::Done(consumed, ());
        } else {
          self.state          = ParserState::Header;
          self.consumer_state = ConsumerState::Continue(consumed);
        }
      }
      IResult::Error(_)      => {
        let kind = ErrorKind::Custom(2);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(_) => {
        let needed = Move::Await(Needed::Size(length as usize));

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }
}

impl<'a> Consumer<&'a [u8], (), ErrorKind, Move> for MetaDataConsumer {
  fn state(&self) -> &ConsumerState<(), ErrorKind, Move> {
    &self.consumer_state
  }

  fn handle(&mut self, input: Input<&'a [u8]>)
            -> &ConsumerState<(), ErrorKind, Move> {
    match input {
      Input::Element(i) | Input::Eof(Some(i)) => {
        match self.state {
          ParserState::Marker      => self.handle_marker(i),
          ParserState::Header      => self.handle_header(i),
          ParserState::Block(data) => self.handle_block(i, data),
        }
      }
      Input::Empty | Input::Eof(None)         => {
        let kind = match self.state {
          ParserState::Marker   => ErrorKind::Custom(0),
          ParserState::Header   => ErrorKind::Custom(1),
          ParserState::Block(_) => ErrorKind::Custom(2),
        };

        self.consumer_state = ConsumerState::Error(kind);
      }
    }

    &self.consumer_state
  }
}
<|endoftext|>


<commit_before>use std::io::Read;
use std::str::FromStr;

use error::Error;
use model::reader::{TmxReader, ElementReader};

define_iterator_wrapper!(Properties, Property);

#[derive(Debug, Default)]
pub struct Property {
    name: String,
    value: String,
    property_type: PropertyType,
}

impl Property {
    pub fn new<S>(name: S, value: S, property_type: PropertyType) -> Property
        where S: Into<String>
    {
        Property {
            name: name.into(),
            value: value.into(),
            property_type: property_type,
        }
    }

    pub fn name(&self) -> &str {
        &self.name
    }

    fn set_name<S: Into<String>>(&mut self, name: S) {
        self.name = name.into();
    }

    pub fn property_type(&self) -> PropertyType {
        self.property_type
    }

    fn set_property_type(&mut self, property_type: PropertyType) {
        self.property_type = property_type;
    }

    pub fn value(&self) -> &str {
        &self.value
    }

    fn set_value<S: Into<String>>(&mut self, value: S) {
        self.value = value.into();
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PropertyType {
    String,
    Int,
    Float,
    Bool,
}

impl Default for PropertyType {
    fn default() -> PropertyType {
        PropertyType::String
    }
}

#[derive(Debug, Default)]
pub struct PropertyCollection(Vec<Property>);

impl PropertyCollection {
    pub fn new() -> PropertyCollection {
        PropertyCollection(Vec::new())
    }

    pub fn push(&mut self, property: Property) {
        self.0.push(property);
    }

    pub fn iter(&self) -> Properties {
        Properties(self.0.iter())
    }
}

impl FromStr for PropertyType {
    type Err = Error;

    fn from_str(s: &str) -> ::Result<PropertyType> {
        match s {
            "string" => Ok(PropertyType::String),
            "int" => Ok(PropertyType::Int),
            "float" => Ok(PropertyType::Float),
            "bool" => Ok(PropertyType::Bool),
            _ => Err(Error::BadPropertyType(s.to_string())),
        }
    }
}

impl<R: Read> ElementReader<Property> for TmxReader<R> {
    fn read_attributes(&mut self, property: &mut Property, name: &str, value: &str) -> ::Result<()> {
        match name {
            "name" => {
                property.set_name(value);
            }
            "type" => {
                property.set_property_type(try!(PropertyType::from_str(value)));
            }
            "value" => {
                property.set_value(value);
            }
            _ => {
                return Err(Error::UnknownAttribute(name.to_string()));
            }
        };
        Ok(())
    }
}

<commit_msg>Added color and file property types<commit_after>use std::io::Read;
use std::str::FromStr;

use error::Error;
use model::reader::{TmxReader, ElementReader};

define_iterator_wrapper!(Properties, Property);

#[derive(Debug, Default)]
pub struct Property {
    name: String,
    value: String,
    property_type: PropertyType,
}

impl Property {
    pub fn new<S>(name: S, value: S, property_type: PropertyType) -> Property
        where S: Into<String>
    {
        Property {
            name: name.into(),
            value: value.into(),
            property_type: property_type,
        }
    }

    pub fn name(&self) -> &str {
        &self.name
    }

    fn set_name<S: Into<String>>(&mut self, name: S) {
        self.name = name.into();
    }

    pub fn property_type(&self) -> PropertyType {
        self.property_type
    }

    fn set_property_type(&mut self, property_type: PropertyType) {
        self.property_type = property_type;
    }

    pub fn value(&self) -> &str {
        &self.value
    }

    fn set_value<S: Into<String>>(&mut self, value: S) {
        self.value = value.into();
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PropertyType {
    Bool,
    Color,
    File,
    Float,
    Int,
    String,
}

impl Default for PropertyType {
    fn default() -> PropertyType {
        PropertyType::String
    }
}

#[derive(Debug, Default)]
pub struct PropertyCollection(Vec<Property>);

impl PropertyCollection {
    pub fn new() -> PropertyCollection {
        PropertyCollection(Vec::new())
    }

    pub fn push(&mut self, property: Property) {
        self.0.push(property);
    }

    pub fn iter(&self) -> Properties {
        Properties(self.0.iter())
    }
}

impl FromStr for PropertyType {
    type Err = Error;

    fn from_str(s: &str) -> ::Result<PropertyType> {
        match s {
            "bool" => Ok(PropertyType::Bool),
            "color" => Ok(PropertyType::Color),
            "file" => Ok(PropertyType::File),
            "float" => Ok(PropertyType::Float),
            "int" => Ok(PropertyType::Int),
            "string" => Ok(PropertyType::String),
            _ => Err(Error::BadPropertyType(s.to_string())),
        }
    }
}

impl<R: Read> ElementReader<Property> for TmxReader<R> {
    fn read_attributes(&mut self, property: &mut Property, name: &str, value: &str) -> ::Result<()> {
        match name {
            "name" => {
                property.set_name(value);
            }
            "type" => {
                property.set_property_type(try!(PropertyType::from_str(value)));
            }
            "value" => {
                property.set_value(value);
            }
            _ => {
                return Err(Error::UnknownAttribute(name.to_string()));
            }
        };
        Ok(())
    }
}

<|endoftext|>


<commit_before>//! Defines the `MoveGenerator` trait.

use std::mem::uninitialized;
use std::cmp::max;
use uci::SetOption;
use board::*;
use moves::*;
use value::*;
use evaluator::Evaluator;
use utils::bitsets::*;
use utils::BoardGeometry;

/// A trait for move generators.
///
/// A `MoveGenerator` holds a chess position and can:
///
/// * Generate all legal moves, or a subset of all legal moves in the
///   current position.
///
/// * Perform static exchange evaluation for the generated moves.
///
/// * Play a selected move and take it back.
///
/// * Provide a static evaluator bound to the current position.
///
/// * Calculate Zobrist hashes.
///
/// **Important note:** `MoveGenerator` is unaware of repeating
/// positions and rule-50.
pub trait MoveGenerator: Sized + Send + Clone + SetOption {
    /// The type of static evaluator that the implementation works
    /// with.
    type Evaluator: Evaluator;

    /// Creates a new instance, consuming the supplied `Board`
    /// instance.
    ///
    /// Returns `None` if the position is illegal.
    fn from_board(board: Board) -> Option<Self>;

    /// Returns the Zobrist hash value for the underlying `Board`
    /// instance.
    ///
    /// Zobrist hashing is a technique to transform a board position
    /// into a number of a fixed length, with an equal distribution
    /// over all possible numbers, invented by Albert Zobrist. The key
    /// property of this method is that two similar positions generate
    /// entirely different hash numbers.
    ///
    /// **Important note:** This method will be relatively slow if the
    /// implementation calculates the hash value "from
    /// scratch". Inspect the implementation before using `hash` in
    /// time-critical paths. (See `do_move`.)
    fn hash(&self) -> u64;

    /// Returns a reference to the underlying `Board` instance.
    #[inline(always)]
    fn board(&self) -> &Board;

    /// Returns a bitboard with all pieces of color `us` that attack
    /// `square`.
    fn attacks_to(&self, us: Color, square: Square) -> Bitboard;

    /// Returns a bitboard with all enemy pieces that attack the king.
    ///
    /// **Important note:** The bitboard of checkers is calculated on
    /// the first call to `checkers`, and is stored in case another
    /// call is made before doing/undoing any moves. In that case
    /// `checkers` returns the saved bitboard instead of
    /// re-calculating it, thus saving time.
    fn checkers(&self) -> Bitboard;

    /// Returns a reference to a static evaluator bound to the current
    /// position.
    fn evaluator(&self) -> &Self::Evaluator;

    /// Generates all legal moves, possibly including some
    /// pseudo-legal moves too.
    ///
    /// The moves are added to `moves`. All generated moves with
    /// pieces other than the king will be legal. Some of the
    /// generated king's moves may be illegal because the destination
    /// square is under attack. This arrangement has two important
    /// advantages:
    ///
    /// * `do_move` can do its work without knowing the set of
    ///   checkers and pinned pieces, so there is no need to keep
    ///   those around.
    ///
    /// * A beta cut-off may make the verification that king's
    ///   destination square is not under attack unnecessary, thus
    ///   saving time.
    ///
    /// **Note:** A pseudo-legal move is a move that is otherwise
    /// legal, except it might leave the king in check.
    fn generate_all<T: AddMove>(&self, moves: &mut T);

    /// Generates moves for the quiescence search.
    ///
    /// The moves are added to `moves`. This method always generates a
    /// **subset** of the moves generated by `generate_all`:
    ///
    /// * If the king is in check, all legal moves are included.
    ///
    /// * Captures and pawn promotions to queen are always included.
    ///
    /// * If `generate_checks` is `true`, moves that give check are
    ///   included too. Discovered checks and checks given by castling
    ///   can be omitted for speed.
    fn generate_forcing<T: AddMove>(&self, generate_checks: bool, moves: &mut T);

    /// Checks if `move_digest` represents a pseudo-legal move.
    ///
    /// If a move `m` exists that would be generated by
    /// `generate_all` if called for the current position on the
    /// board, and for that move `m.digest() == move_digest`, this
    /// method will return `Some(m)`. Otherwise it will return
    /// `None`. This is useful when playing moves from the
    /// transposition table, without calling `generate_all`.
    fn try_move_digest(&self, move_digest: MoveDigest) -> Option<Move>;

    /// Returns a null move.
    ///
    /// "Null move" is a pseudo-move that changes only the side to
    /// move. It is sometimes useful to include a speculative null
    /// move in the search tree so as to achieve more aggressive
    /// pruning. Null moves are represented as king's moves for which
    /// the origin and destination squares are the same.
    fn null_move(&self) -> Move;

    /// Plays a move on the board.
    ///
    /// It verifies if the move is legal. If the move is legal, the
    /// board is updated and an `u64` value is returned, which should
    /// be XOR-ed with old board's hash value to obtain new board's
    /// hash value. If the move is illegal, `None` is returned without
    /// updating the board. The move passed to this method **must**
    /// have been generated by `generate_all`, `generate_forcing`,
    /// `try_move_digest`, or `null_move` methods for the current
    /// position on the board.
    ///
    /// The moves generated by the `null_move` method are
    /// exceptions. For them `do_move` will return `None` if and only
    /// if the king is in check.
    fn do_move(&mut self, m: Move) -> Option<u64>;

    /// Takes back last played move.
    ///
    /// The move passed to this method **must** be the last move passed
    /// to `do_move`.
    fn undo_move(&mut self, m: Move);

    /// Returns the likely evaluation change (material) to be lost or
    /// gained as a result of a given move.
    ///
    /// This method performs static exchange evaluation (SEE). It
    /// examines the consequence of a series of exchanges on the
    /// destination square after a given move. A positive returned
    /// value indicates a "winning" move. For example, "PxQ" will
    /// always be a win, since the pawn side can choose to stop the
    /// exchange after its pawn is recaptured, and still be ahead. SEE
    /// is just an evaluation calculated without actually trying moves
    /// on the board, and therefore the returned value might be
    /// incorrect.
    ///
    /// The move passed to this method must have been generated by
    /// `generate_all`, `generate_forcing`, `try_move_digest`, or
    /// `null_move` methods for the current position on the board.
    fn evaluate_move(&self, m: Move) -> Value {
        debug_assert!(m.played_piece() < PIECE_NONE);
        debug_assert!(m.captured_piece() <= PIECE_NONE);
        const PIECE_VALUES: [Value; 8] = [10000, 975, 500, 325, 325, 100, 0, 0];

        // This is the square on which all the action takes place.
        let exchange_square = m.dest_square();

        unsafe {
            let geometry = BoardGeometry::get();
            let behind_blocker: &[Bitboard; 64] = geometry.squares_behind_blocker
                                                          .get_unchecked(exchange_square);
            let piece_type: &[Bitboard; 6] = &self.board().pieces.piece_type;
            let color: &[Bitboard; 2] = &self.board().pieces.color;
            let straight_sliders = piece_type[QUEEN] | piece_type[ROOK];
            let diag_sliders = piece_type[QUEEN] | piece_type[BISHOP];

            // `may_xray` holds the set of pieces that may block attacks
            // from other pieces. We will consider adding new
            // attackers/defenders every time a piece from the `may_xray`
            // set makes a capture.
            let may_xray = piece_type[PAWN] | piece_type[BISHOP] | piece_type[ROOK] |
                           piece_type[QUEEN];

            // These variables will be updated on each capture:
            let mut us = self.board().to_move;
            let mut depth = 0;
            let mut piece;
            let mut orig_square_bb = 1 << m.orig_square();
            let mut attackers_and_defenders = self.attacks_to(WHITE, exchange_square) |
                                              self.attacks_to(BLACK, exchange_square);

            // The `gain` array will hold the total material gained at
            // each `depth`, from the viewpoint of the side that made the
            // last capture (`us`).
            let mut gain: [Value; 34] = uninitialized();
            gain[0] = if m.move_type() == MOVE_PROMOTION {
                piece = Move::piece_from_aux_data(m.aux_data());
                PIECE_VALUES[m.captured_piece()] + PIECE_VALUES[piece] - PIECE_VALUES[PAWN]
            } else {
                piece = m.played_piece();
                *PIECE_VALUES.get_unchecked(m.captured_piece())
            };

            // Examine the possible exchanges, fill the `gain` array.
            'exchange: while orig_square_bb != 0 {
                let current_gain = *gain.get_unchecked(depth);

                // Store a speculative value that will be used if the
                // captured piece happens to be defended.
                let speculative_gain: &mut Value = gain.get_unchecked_mut(depth + 1);
                *speculative_gain = *PIECE_VALUES.get_unchecked(piece) - current_gain;

                if max(-current_gain, *speculative_gain) < 0 {
                    // The side that made the last capture wins even if
                    // the captured piece happens to be defended. So, we
                    // stop here to save precious CPU cycles. Note that
                    // here we may happen to return an incorrect SEE
                    // value, but the sign will be correct, which is by
                    // far the most important information.
                    break;
                }

                // Register that `orig_square_bb` is now vacant.
                attackers_and_defenders &= !orig_square_bb;

                // Consider adding new attackers/defenders, now that
                // `orig_square_bb` is vacant.
                if orig_square_bb & may_xray != 0 {
                    attackers_and_defenders |= {
                        let behind = self.board().occupied &
                                     *behind_blocker.get_unchecked(bitscan_forward(orig_square_bb));
                        match behind & straight_sliders &
                              geometry.attacks_from_unsafe(ROOK, exchange_square, behind) {
                            0 => {
                                // Not a straight slider -- possibly a diagonal slider.
                                behind & diag_sliders &
                                geometry.attacks_from_unsafe(BISHOP, exchange_square, behind)
                            }
                            x => x,
                        }
                    };
                }

                // Change the side to move.
                us ^= 1;

                // Find the next piece to enter the exchange. (The least
                // valuable piece belonging to the side to move.)
                let candidates = attackers_and_defenders & *color.get_unchecked(us);
                if candidates != 0 {
                    for p in (KING..PIECE_NONE).rev() {
                        let bb = candidates & piece_type[p];
                        if bb != 0 {
                            depth += 1;
                            piece = p;
                            orig_square_bb = ls1b(bb);
                            continue 'exchange;
                        }
                    }
                }
                break 'exchange;
            }

            // Negamax the `gain` array for the final static exchange
            // evaluation. (The `gain` array actually represents an unary
            // tree, at each node of which the player can either continue
            // the exchange or back off.)
            while depth > 0 {
                *gain.get_unchecked_mut(depth - 1) = -max(-*gain.get_unchecked(depth - 1),
                                                          *gain.get_unchecked(depth));
                depth -= 1;
            }
            gain[0]
        }
    }
}
<commit_msg>optimize evaluate_move()<commit_after>//! Defines the `MoveGenerator` trait.

use std::mem::uninitialized;
use std::cmp::max;
use uci::SetOption;
use board::*;
use moves::*;
use value::*;
use evaluator::Evaluator;
use utils::bitsets::*;
use utils::BoardGeometry;

/// A trait for move generators.
///
/// A `MoveGenerator` holds a chess position and can:
///
/// * Generate all legal moves, or a subset of all legal moves in the
///   current position.
///
/// * Perform static exchange evaluation for the generated moves.
///
/// * Play a selected move and take it back.
///
/// * Provide a static evaluator bound to the current position.
///
/// * Calculate Zobrist hashes.
///
/// **Important note:** `MoveGenerator` is unaware of repeating
/// positions and rule-50.
pub trait MoveGenerator: Sized + Send + Clone + SetOption {
    /// The type of static evaluator that the implementation works
    /// with.
    type Evaluator: Evaluator;

    /// Creates a new instance, consuming the supplied `Board`
    /// instance.
    ///
    /// Returns `None` if the position is illegal.
    fn from_board(board: Board) -> Option<Self>;

    /// Returns the Zobrist hash value for the underlying `Board`
    /// instance.
    ///
    /// Zobrist hashing is a technique to transform a board position
    /// into a number of a fixed length, with an equal distribution
    /// over all possible numbers, invented by Albert Zobrist. The key
    /// property of this method is that two similar positions generate
    /// entirely different hash numbers.
    ///
    /// **Important note:** This method will be relatively slow if the
    /// implementation calculates the hash value "from
    /// scratch". Inspect the implementation before using `hash` in
    /// time-critical paths. (See `do_move`.)
    fn hash(&self) -> u64;

    /// Returns a reference to the underlying `Board` instance.
    #[inline(always)]
    fn board(&self) -> &Board;

    /// Returns a bitboard with all pieces of color `us` that attack
    /// `square`.
    fn attacks_to(&self, us: Color, square: Square) -> Bitboard;

    /// Returns a bitboard with all enemy pieces that attack the king.
    ///
    /// **Important note:** The bitboard of checkers is calculated on
    /// the first call to `checkers`, and is stored in case another
    /// call is made before doing/undoing any moves. In that case
    /// `checkers` returns the saved bitboard instead of
    /// re-calculating it, thus saving time.
    fn checkers(&self) -> Bitboard;

    /// Returns a reference to a static evaluator bound to the current
    /// position.
    fn evaluator(&self) -> &Self::Evaluator;

    /// Generates all legal moves, possibly including some
    /// pseudo-legal moves too.
    ///
    /// The moves are added to `moves`. All generated moves with
    /// pieces other than the king will be legal. Some of the
    /// generated king's moves may be illegal because the destination
    /// square is under attack. This arrangement has two important
    /// advantages:
    ///
    /// * `do_move` can do its work without knowing the set of
    ///   checkers and pinned pieces, so there is no need to keep
    ///   those around.
    ///
    /// * A beta cut-off may make the verification that king's
    ///   destination square is not under attack unnecessary, thus
    ///   saving time.
    ///
    /// **Note:** A pseudo-legal move is a move that is otherwise
    /// legal, except it might leave the king in check.
    fn generate_all<T: AddMove>(&self, moves: &mut T);

    /// Generates moves for the quiescence search.
    ///
    /// The moves are added to `moves`. This method always generates a
    /// **subset** of the moves generated by `generate_all`:
    ///
    /// * If the king is in check, all legal moves are included.
    ///
    /// * Captures and pawn promotions to queen are always included.
    ///
    /// * If `generate_checks` is `true`, moves that give check are
    ///   included too. Discovered checks and checks given by castling
    ///   can be omitted for speed.
    fn generate_forcing<T: AddMove>(&self, generate_checks: bool, moves: &mut T);

    /// Checks if `move_digest` represents a pseudo-legal move.
    ///
    /// If a move `m` exists that would be generated by
    /// `generate_all` if called for the current position on the
    /// board, and for that move `m.digest() == move_digest`, this
    /// method will return `Some(m)`. Otherwise it will return
    /// `None`. This is useful when playing moves from the
    /// transposition table, without calling `generate_all`.
    fn try_move_digest(&self, move_digest: MoveDigest) -> Option<Move>;

    /// Returns a null move.
    ///
    /// "Null move" is a pseudo-move that changes only the side to
    /// move. It is sometimes useful to include a speculative null
    /// move in the search tree so as to achieve more aggressive
    /// pruning. Null moves are represented as king's moves for which
    /// the origin and destination squares are the same.
    fn null_move(&self) -> Move;

    /// Plays a move on the board.
    ///
    /// It verifies if the move is legal. If the move is legal, the
    /// board is updated and an `u64` value is returned, which should
    /// be XOR-ed with old board's hash value to obtain new board's
    /// hash value. If the move is illegal, `None` is returned without
    /// updating the board. The move passed to this method **must**
    /// have been generated by `generate_all`, `generate_forcing`,
    /// `try_move_digest`, or `null_move` methods for the current
    /// position on the board.
    ///
    /// The moves generated by the `null_move` method are
    /// exceptions. For them `do_move` will return `None` if and only
    /// if the king is in check.
    fn do_move(&mut self, m: Move) -> Option<u64>;

    /// Takes back last played move.
    ///
    /// The move passed to this method **must** be the last move passed
    /// to `do_move`.
    fn undo_move(&mut self, m: Move);

    /// Returns the likely evaluation change (material) to be lost or
    /// gained as a result of a given move.
    ///
    /// This method performs static exchange evaluation (SEE). It
    /// examines the consequence of a series of exchanges on the
    /// destination square after a given move. A positive returned
    /// value indicates a "winning" move. For example, "PxQ" will
    /// always be a win, since the pawn side can choose to stop the
    /// exchange after its pawn is recaptured, and still be ahead. SEE
    /// is just an evaluation calculated without actually trying moves
    /// on the board, and therefore the returned value might be
    /// incorrect.
    ///
    /// The move passed to this method must have been generated by
    /// `generate_all`, `generate_forcing`, `try_move_digest`, or
    /// `null_move` methods for the current position on the board.
    fn evaluate_move(&self, m: Move) -> Value {
        debug_assert!(m.played_piece() < PIECE_NONE);
        debug_assert!(m.captured_piece() <= PIECE_NONE);
        const PIECE_VALUES: [Value; 8] = [10000, 975, 500, 325, 325, 100, 0, 0];

        // This is the square on which all the action takes place.
        let exchange_square = m.dest_square();

        unsafe {
            let geometry = BoardGeometry::get();
            let behind_blocker: &[Bitboard; 64] = geometry.squares_behind_blocker
                                                          .get_unchecked(exchange_square);
            let piece_type: &[Bitboard; 6] = &self.board().pieces.piece_type;
            let color: &[Bitboard; 2] = &self.board().pieces.color;
            let straight_sliders = piece_type[QUEEN] | piece_type[ROOK];
            let diag_sliders = piece_type[QUEEN] | piece_type[BISHOP];

            // `may_xray` holds the set of pieces that may block attacks
            // from other pieces. We will consider adding new
            // attackers/defenders every time a piece from the `may_xray`
            // set makes a capture.
            let may_xray = piece_type[PAWN] | piece_type[BISHOP] | piece_type[ROOK] |
                           piece_type[QUEEN];

            // These variables will be updated on each capture:
            let mut us = self.board().to_move;
            let mut depth = 0;
            let mut piece;
            let mut orig_square_bb = 1 << m.orig_square();
            let mut attackers_and_defenders = self.attacks_to(WHITE, exchange_square) |
                                              self.attacks_to(BLACK, exchange_square);

            // The `gain` array will hold the total material gained at
            // each `depth`, from the viewpoint of the side that made the
            // last capture (`us`).
            let mut gain: [Value; 34] = uninitialized();
            gain[0] = if m.move_type() == MOVE_PROMOTION {
                piece = Move::piece_from_aux_data(m.aux_data());
                PIECE_VALUES[m.captured_piece()] + PIECE_VALUES[piece] - PIECE_VALUES[PAWN]
            } else {
                piece = m.played_piece();
                *PIECE_VALUES.get_unchecked(m.captured_piece())
            };

            // Examine the possible exchanges, fill the `gain` array.
            'exchange: while orig_square_bb != 0 {
                let current_gain = *gain.get_unchecked(depth);

                // Store a speculative value that will be used if the
                // captured piece happens to be defended.
                let speculative_gain: &mut Value = gain.get_unchecked_mut(depth + 1);
                *speculative_gain = *PIECE_VALUES.get_unchecked(piece) - current_gain;

                if max(-current_gain, *speculative_gain) < 0 {
                    // The side that made the last capture wins even if
                    // the captured piece happens to be defended. So, we
                    // stop here to save precious CPU cycles. Note that
                    // here we may happen to return an incorrect SEE
                    // value, but the sign will be correct, which is by
                    // far the most important information.
                    break;
                }

                // Register that `orig_square_bb` is now vacant.
                attackers_and_defenders &= !orig_square_bb;

                // Consider adding new attackers/defenders, now that
                // `orig_square_bb` is vacant.
                if orig_square_bb & may_xray != 0 {
                    let behind = self.board().occupied &
                                 *behind_blocker.get_unchecked(bitscan_forward(orig_square_bb));
                    if behind & (straight_sliders | diag_sliders) != 0 {
                        attackers_and_defenders |=
                            match behind & straight_sliders &
                                  geometry.attacks_from_unsafe(ROOK, exchange_square, behind) {
                                0 => {
                                    // Not a straight slider, possibly a diagonal slider.
                                    behind & diag_sliders &
                                    geometry.attacks_from_unsafe(BISHOP, exchange_square, behind)
                                }
                                bb => {
                                    // A straight slider.
                                    bb
                                }
                            };
                    }
                }

                // Change the side to move.
                us ^= 1;

                // Find the next piece to enter the exchange. (The least
                // valuable piece belonging to the side to move.)
                let candidates = attackers_and_defenders & *color.get_unchecked(us);
                if candidates != 0 {
                    for p in (KING..PIECE_NONE).rev() {
                        let bb = candidates & piece_type[p];
                        if bb != 0 {
                            depth += 1;
                            piece = p;
                            orig_square_bb = ls1b(bb);
                            continue 'exchange;
                        }
                    }
                }
                break 'exchange;
            }

            // Negamax the `gain` array for the final static exchange
            // evaluation. (The `gain` array actually represents an unary
            // tree, at each node of which the player can either continue
            // the exchange or back off.)
            while depth > 0 {
                *gain.get_unchecked_mut(depth - 1) = -max(-*gain.get_unchecked(depth - 1),
                                                          *gain.get_unchecked(depth));
                depth -= 1;
            }
            gain[0]
        }
    }
}
<|endoftext|>


<commit_before>//! The underlying OsString/OsStr implementation on Unix and many other
//! systems: just a `Vec<u8>`/`[u8]`.

use crate::borrow::Cow;

use crate::fmt;
use crate::mem;
use crate::rc::Rc;
use crate::str;
use crate::sync::Arc;
use crate::sys_common::bytestring::debug_fmt_bytestring;
use crate::sys_common::{AsInner, IntoInner};

use core::str::lossy::Utf8Lossy;

#[derive(Hash)]
pub(crate) struct Buf {
    pub inner: Vec<u8>,
}

// FIXME:
// `Buf::as_slice` current implementation relies
// on `Slice` being layout-compatible with `[u8]`.
// When attribute privacy is implemented, `Slice` should be annotated as `#[repr(transparent)]`.
// Anyway, `Slice` representation and layout are considered implementation detail, are
// not documented and must not be relied upon.
pub(crate) struct Slice {
    pub inner: [u8],
}

impl fmt::Debug for Slice {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        debug_fmt_bytestring(&self.inner, formatter)
    }
}

impl fmt::Display for Slice {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&Utf8Lossy::from_bytes(&self.inner), formatter)
    }
}

impl fmt::Debug for Buf {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(self.as_slice(), formatter)
    }
}

impl fmt::Display for Buf {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(self.as_slice(), formatter)
    }
}

impl Clone for Buf {
    #[inline]
    fn clone(&self) -> Self {
        Buf { inner: self.inner.clone() }
    }

    #[inline]
    fn clone_from(&mut self, source: &Self) {
        self.inner.clone_from(&source.inner)
    }
}

impl IntoInner<Vec<u8>> for Buf {
    fn into_inner(self) -> Vec<u8> {
        self.inner
    }
}

impl AsInner<[u8]> for Buf {
    fn as_inner(&self) -> &[u8] {
        &self.inner
    }
}

impl Buf {
    pub fn from_string(s: String) -> Buf {
        Buf { inner: s.into_bytes() }
    }

    #[inline]
    pub fn with_capacity(capacity: usize) -> Buf {
        Buf { inner: Vec::with_capacity(capacity) }
    }

    #[inline]
    pub fn clear(&mut self) {
        self.inner.clear()
    }

    #[inline]
    pub fn capacity(&self) -> usize {
        self.inner.capacity()
    }

    #[inline]
    pub fn reserve(&mut self, additional: usize) {
        self.inner.reserve(additional)
    }

    #[inline]
    pub fn reserve_exact(&mut self, additional: usize) {
        self.inner.reserve_exact(additional)
    }

    #[inline]
    pub fn shrink_to_fit(&mut self) {
        self.inner.shrink_to_fit()
    }

    #[inline]
    pub fn shrink_to(&mut self, min_capacity: usize) {
        self.inner.shrink_to(min_capacity)
    }

    #[inline]
    pub fn as_slice(&self) -> &Slice {
        // SAFETY: Slice just wraps [u8],
        // and &*self.inner is &[u8], therefore
        // transmuting &[u8] to &Slice is safe.
        unsafe { mem::transmute(&*self.inner) }
    }

    #[inline]
    pub fn as_mut_slice(&mut self) -> &mut Slice {
        // SAFETY: Slice just wraps [u8],
        // and &mut *self.inner is &mut [u8], therefore
        // transmuting &mut [u8] to &mut Slice is safe.
        unsafe { mem::transmute(&mut *self.inner) }
    }

    pub fn into_string(self) -> Result<String, Buf> {
        String::from_utf8(self.inner).map_err(|p| Buf { inner: p.into_bytes() })
    }

    pub fn push_slice(&mut self, s: &Slice) {
        self.inner.extend_from_slice(&s.inner)
    }

    #[inline]
    pub fn into_box(self) -> Box<Slice> {
        unsafe { mem::transmute(self.inner.into_boxed_slice()) }
    }

    #[inline]
    pub fn from_box(boxed: Box<Slice>) -> Buf {
        let inner: Box<[u8]> = unsafe { mem::transmute(boxed) };
        Buf { inner: inner.into_vec() }
    }

    #[inline]
    pub fn into_arc(&self) -> Arc<Slice> {
        self.as_slice().into_arc()
    }

    #[inline]
    pub fn into_rc(&self) -> Rc<Slice> {
        self.as_slice().into_rc()
    }
}

impl Slice {
    #[inline]
    fn from_u8_slice(s: &[u8]) -> &Slice {
        unsafe { mem::transmute(s) }
    }

    #[inline]
    pub fn from_str(s: &str) -> &Slice {
        Slice::from_u8_slice(s.as_bytes())
    }

    pub fn to_str(&self) -> Option<&str> {
        str::from_utf8(&self.inner).ok()
    }

    pub fn to_string_lossy(&self) -> Cow<'_, str> {
        String::from_utf8_lossy(&self.inner)
    }

    pub fn to_owned(&self) -> Buf {
        Buf { inner: self.inner.to_vec() }
    }

    pub fn clone_into(&self, buf: &mut Buf) {
        self.inner.clone_into(&mut buf.inner)
    }

    #[inline]
    pub fn into_box(&self) -> Box<Slice> {
        let boxed: Box<[u8]> = self.inner.into();
        unsafe { mem::transmute(boxed) }
    }

    pub fn empty_box() -> Box<Slice> {
        let boxed: Box<[u8]> = Default::default();
        unsafe { mem::transmute(boxed) }
    }

    #[inline]
    pub fn into_arc(&self) -> Arc<Slice> {
        let arc: Arc<[u8]> = Arc::from(&self.inner);
        unsafe { Arc::from_raw(Arc::into_raw(arc) as *const Slice) }
    }

    #[inline]
    pub fn into_rc(&self) -> Rc<Slice> {
        let rc: Rc<[u8]> = Rc::from(&self.inner);
        unsafe { Rc::from_raw(Rc::into_raw(rc) as *const Slice) }
    }

    #[inline]
    pub fn make_ascii_lowercase(&mut self) {
        self.inner.make_ascii_lowercase()
    }

    #[inline]
    pub fn make_ascii_uppercase(&mut self) {
        self.inner.make_ascii_uppercase()
    }

    #[inline]
    pub fn to_ascii_lowercase(&self) -> Buf {
        Buf { inner: self.inner.to_ascii_lowercase() }
    }

    #[inline]
    pub fn to_ascii_uppercase(&self) -> Buf {
        Buf { inner: self.inner.to_ascii_uppercase() }
    }

    #[inline]
    pub fn is_ascii(&self) -> bool {
        self.inner.is_ascii()
    }

    #[inline]
    pub fn eq_ignore_ascii_case(&self, other: &Self) -> bool {
        self.inner.eq_ignore_ascii_case(&other.inner)
    }
}
<commit_msg>Make `os_str_bytes::{Buf, Slice}` `pub` and `repr(transparent)`<commit_after>//! The underlying OsString/OsStr implementation on Unix and many other
//! systems: just a `Vec<u8>`/`[u8]`.

use crate::borrow::Cow;

use crate::fmt;
use crate::mem;
use crate::rc::Rc;
use crate::str;
use crate::sync::Arc;
use crate::sys_common::bytestring::debug_fmt_bytestring;
use crate::sys_common::{AsInner, IntoInner};

use core::str::lossy::Utf8Lossy;

#[derive(Hash)]
#[repr(transparent)]
pub struct Buf {
    pub inner: Vec<u8>,
}

#[repr(transparent)]
pub struct Slice {
    pub inner: [u8],
}

impl fmt::Debug for Slice {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        debug_fmt_bytestring(&self.inner, formatter)
    }
}

impl fmt::Display for Slice {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&Utf8Lossy::from_bytes(&self.inner), formatter)
    }
}

impl fmt::Debug for Buf {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(self.as_slice(), formatter)
    }
}

impl fmt::Display for Buf {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(self.as_slice(), formatter)
    }
}

impl Clone for Buf {
    #[inline]
    fn clone(&self) -> Self {
        Buf { inner: self.inner.clone() }
    }

    #[inline]
    fn clone_from(&mut self, source: &Self) {
        self.inner.clone_from(&source.inner)
    }
}

impl IntoInner<Vec<u8>> for Buf {
    fn into_inner(self) -> Vec<u8> {
        self.inner
    }
}

impl AsInner<[u8]> for Buf {
    fn as_inner(&self) -> &[u8] {
        &self.inner
    }
}

impl Buf {
    pub fn from_string(s: String) -> Buf {
        Buf { inner: s.into_bytes() }
    }

    #[inline]
    pub fn with_capacity(capacity: usize) -> Buf {
        Buf { inner: Vec::with_capacity(capacity) }
    }

    #[inline]
    pub fn clear(&mut self) {
        self.inner.clear()
    }

    #[inline]
    pub fn capacity(&self) -> usize {
        self.inner.capacity()
    }

    #[inline]
    pub fn reserve(&mut self, additional: usize) {
        self.inner.reserve(additional)
    }

    #[inline]
    pub fn reserve_exact(&mut self, additional: usize) {
        self.inner.reserve_exact(additional)
    }

    #[inline]
    pub fn shrink_to_fit(&mut self) {
        self.inner.shrink_to_fit()
    }

    #[inline]
    pub fn shrink_to(&mut self, min_capacity: usize) {
        self.inner.shrink_to(min_capacity)
    }

    #[inline]
    pub fn as_slice(&self) -> &Slice {
        // SAFETY: Slice just wraps [u8],
        // and &*self.inner is &[u8], therefore
        // transmuting &[u8] to &Slice is safe.
        unsafe { mem::transmute(&*self.inner) }
    }

    #[inline]
    pub fn as_mut_slice(&mut self) -> &mut Slice {
        // SAFETY: Slice just wraps [u8],
        // and &mut *self.inner is &mut [u8], therefore
        // transmuting &mut [u8] to &mut Slice is safe.
        unsafe { mem::transmute(&mut *self.inner) }
    }

    pub fn into_string(self) -> Result<String, Buf> {
        String::from_utf8(self.inner).map_err(|p| Buf { inner: p.into_bytes() })
    }

    pub fn push_slice(&mut self, s: &Slice) {
        self.inner.extend_from_slice(&s.inner)
    }

    #[inline]
    pub fn into_box(self) -> Box<Slice> {
        unsafe { mem::transmute(self.inner.into_boxed_slice()) }
    }

    #[inline]
    pub fn from_box(boxed: Box<Slice>) -> Buf {
        let inner: Box<[u8]> = unsafe { mem::transmute(boxed) };
        Buf { inner: inner.into_vec() }
    }

    #[inline]
    pub fn into_arc(&self) -> Arc<Slice> {
        self.as_slice().into_arc()
    }

    #[inline]
    pub fn into_rc(&self) -> Rc<Slice> {
        self.as_slice().into_rc()
    }
}

impl Slice {
    #[inline]
    fn from_u8_slice(s: &[u8]) -> &Slice {
        unsafe { mem::transmute(s) }
    }

    #[inline]
    pub fn from_str(s: &str) -> &Slice {
        Slice::from_u8_slice(s.as_bytes())
    }

    pub fn to_str(&self) -> Option<&str> {
        str::from_utf8(&self.inner).ok()
    }

    pub fn to_string_lossy(&self) -> Cow<'_, str> {
        String::from_utf8_lossy(&self.inner)
    }

    pub fn to_owned(&self) -> Buf {
        Buf { inner: self.inner.to_vec() }
    }

    pub fn clone_into(&self, buf: &mut Buf) {
        self.inner.clone_into(&mut buf.inner)
    }

    #[inline]
    pub fn into_box(&self) -> Box<Slice> {
        let boxed: Box<[u8]> = self.inner.into();
        unsafe { mem::transmute(boxed) }
    }

    pub fn empty_box() -> Box<Slice> {
        let boxed: Box<[u8]> = Default::default();
        unsafe { mem::transmute(boxed) }
    }

    #[inline]
    pub fn into_arc(&self) -> Arc<Slice> {
        let arc: Arc<[u8]> = Arc::from(&self.inner);
        unsafe { Arc::from_raw(Arc::into_raw(arc) as *const Slice) }
    }

    #[inline]
    pub fn into_rc(&self) -> Rc<Slice> {
        let rc: Rc<[u8]> = Rc::from(&self.inner);
        unsafe { Rc::from_raw(Rc::into_raw(rc) as *const Slice) }
    }

    #[inline]
    pub fn make_ascii_lowercase(&mut self) {
        self.inner.make_ascii_lowercase()
    }

    #[inline]
    pub fn make_ascii_uppercase(&mut self) {
        self.inner.make_ascii_uppercase()
    }

    #[inline]
    pub fn to_ascii_lowercase(&self) -> Buf {
        Buf { inner: self.inner.to_ascii_lowercase() }
    }

    #[inline]
    pub fn to_ascii_uppercase(&self) -> Buf {
        Buf { inner: self.inner.to_ascii_uppercase() }
    }

    #[inline]
    pub fn is_ascii(&self) -> bool {
        self.inner.is_ascii()
    }

    #[inline]
    pub fn eq_ignore_ascii_case(&self, other: &Self) -> bool {
        self.inner.eq_ignore_ascii_case(&other.inner)
    }
}
<|endoftext|>


<commit_before>use crate::core::MANAGED_FILEPATH;
use crate::directory::error::{DeleteError, IOError, LockError, OpenReadError, OpenWriteError};
use crate::directory::DirectoryLock;
use crate::directory::Lock;
use crate::directory::META_LOCK;
use crate::directory::{ReadOnlySource, WritePtr};
use crate::directory::{WatchCallback, WatchHandle};
use crate::error::DataCorruption;
use crate::Directory;
use crate::Result;
use serde_json;
use std::collections::HashSet;
use std::io;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::result;
use std::sync::RwLockWriteGuard;
use std::sync::{Arc, RwLock};

/// Returns true iff the file is "managed".
/// Non-managed file are not subject to garbage collection.
///
/// Filenames that starts by a "." -typically locks-
/// are not managed.
fn is_managed(path: &Path) -> bool {
    path.to_str()
        .map(|p_str| !p_str.starts_with('.'))
        .unwrap_or(true)
}

/// Wrapper of directories that keeps track of files created by Tantivy.
///
/// A managed directory is just a wrapper of a directory
/// that keeps a (persisted) list of the files that
/// have been created (and not deleted) by tantivy so far.
///
/// Thanks to this list, it implements a `garbage_collect` method
/// that removes the files that were created by tantivy and are not
/// useful anymore.
#[derive(Debug)]
pub struct ManagedDirectory {
    directory: Box<dyn Directory>,
    meta_informations: Arc<RwLock<MetaInformation>>,
}

#[derive(Debug, Default)]
struct MetaInformation {
    managed_paths: HashSet<PathBuf>,
}

/// Saves the file containing the list of existing files
/// that were created by tantivy.
fn save_managed_paths(
    directory: &mut dyn Directory,
    wlock: &RwLockWriteGuard<'_, MetaInformation>,
) -> io::Result<()> {
    let mut w = serde_json::to_vec(&wlock.managed_paths)?;
    writeln!(&mut w)?;
    directory.atomic_write(&MANAGED_FILEPATH, &w[..])?;
    Ok(())
}

impl ManagedDirectory {
    /// Wraps a directory as managed directory.
    pub fn wrap<Dir: Directory>(directory: Dir) -> Result<ManagedDirectory> {
        match directory.atomic_read(&MANAGED_FILEPATH) {
            Ok(data) => {
                let managed_files_json = String::from_utf8_lossy(&data);
                let managed_files: HashSet<PathBuf> = serde_json::from_str(&managed_files_json)
                    .map_err(|e| {
                        DataCorruption::new(
                            MANAGED_FILEPATH.to_path_buf(),
                            format!("Managed file cannot be deserialized: {:?}. ", e),
                        )
                    })?;
                Ok(ManagedDirectory {
                    directory: Box::new(directory),
                    meta_informations: Arc::new(RwLock::new(MetaInformation {
                        managed_paths: managed_files,
                    })),
                })
            }
            Err(OpenReadError::FileDoesNotExist(_)) => Ok(ManagedDirectory {
                directory: Box::new(directory),
                meta_informations: Arc::default(),
            }),
            Err(OpenReadError::IOError(e)) => Err(From::from(e)),
        }
    }

    /// Garbage collect unused files.
    ///
    /// Removes the files that were created by `tantivy` and are not
    /// used by any segment anymore.
    ///
    /// * `living_files` - List of files that are still used by the index.
    ///
    /// The use a callback ensures that the list of living_files is computed
    /// while we hold the lock on meta.
    ///
    /// This method does not panick nor returns errors.
    /// If a file cannot be deleted (for permission reasons for instance)
    /// an error is simply logged, and the file remains in the list of managed
    /// files.
    pub fn garbage_collect<L: FnOnce() -> HashSet<PathBuf>>(&mut self, get_living_files: L) {
        info!("Garbage collect");
        let mut files_to_delete = vec![];

        // It is crucial to get the living files after acquiring the
        // read lock of meta informations. That way, we
        // avoid the following scenario.
        //
        // 1) we get the list of living files.
        // 2) someone creates a new file.
        // 3) we start garbage collection and remove this file
        // even though it is a living file.
        //
        // releasing the lock as .delete() will use it too.
        {
            let meta_informations_rlock = self
                .meta_informations
                .read()
                .expect("Managed directory rlock poisoned in garbage collect.");

            // The point of this second "file" lock is to enforce the following scenario
            // 1) process B tries to load a new set of searcher.
            // The list of segments is loaded
            // 2) writer change meta.json (for instance after a merge or a commit)
            // 3) gc kicks in.
            // 4) gc removes a file that was useful for process B, before process B opened it.
            if let Ok(_meta_lock) = self.acquire_lock(&META_LOCK) {
                let living_files = get_living_files();
                for managed_path in &meta_informations_rlock.managed_paths {
                    if !living_files.contains(managed_path) {
                        files_to_delete.push(managed_path.clone());
                    }
                }
            }
        }

        let mut deleted_files = vec![];
        {
            for file_to_delete in files_to_delete {
                match self.delete(&file_to_delete) {
                    Ok(_) => {
                        info!("Deleted {:?}", file_to_delete);
                        deleted_files.push(file_to_delete);
                    }
                    Err(file_error) => {
                        match file_error {
                            DeleteError::FileDoesNotExist(_) => {
                                deleted_files.push(file_to_delete);
                            }
                            DeleteError::IOError(_) => {
                                if !cfg!(target_os = "windows") {
                                    // On windows, delete is expected to fail if the file
                                    // is mmapped.
                                    error!("Failed to delete {:?}", file_to_delete);
                                }
                            }
                        }
                    }
                }
            }
        }

        if !deleted_files.is_empty() {
            // update the list of managed files by removing
            // the file that were removed.
            let mut meta_informations_wlock = self
                .meta_informations
                .write()
                .expect("Managed directory wlock poisoned (2).");
            {
                let managed_paths_write = &mut meta_informations_wlock.managed_paths;
                for delete_file in &deleted_files {
                    managed_paths_write.remove(delete_file);
                }
            }
            if save_managed_paths(self.directory.as_mut(), &meta_informations_wlock).is_err() {
                error!("Failed to save the list of managed files.");
            }
        }
    }

    /// Registers a file as managed
    ///
    /// This method must be called before the file is
    /// actually created to ensure that a failure between
    /// registering the filepath and creating the file
    /// will not lead to garbage files that will
    /// never get removed.
    ///
    /// File starting by "." are reserved to locks.
    /// They are not managed and cannot be subjected
    /// to garbage collection.
    fn register_file_as_managed(&mut self, filepath: &Path) -> io::Result<()> {
        // Files starting by "." (e.g. lock files) are not managed.
        if !is_managed(filepath) {
            return Ok(());
        }
        let mut meta_wlock = self
            .meta_informations
            .write()
            .expect("Managed file lock poisoned");
        let has_changed = meta_wlock.managed_paths.insert(filepath.to_owned());
        if has_changed {
            save_managed_paths(self.directory.as_mut(), &meta_wlock)?;
        }
        Ok(())
    }
}

impl Directory for ManagedDirectory {
    fn open_read(&self, path: &Path) -> result::Result<ReadOnlySource, OpenReadError> {
        self.directory.open_read(path)
    }

    fn open_write(&mut self, path: &Path) -> result::Result<WritePtr, OpenWriteError> {
        self.register_file_as_managed(path)
            .map_err(|e| IOError::with_path(path.to_owned(), e))?;
        self.directory.open_write(path)
    }

    fn atomic_write(&mut self, path: &Path, data: &[u8]) -> io::Result<()> {
        self.register_file_as_managed(path)?;
        self.directory.atomic_write(path, data)
    }

    fn atomic_read(&self, path: &Path) -> result::Result<Vec<u8>, OpenReadError> {
        self.directory.atomic_read(path)
    }

    fn delete(&self, path: &Path) -> result::Result<(), DeleteError> {
        self.directory.delete(path)
    }

    fn exists(&self, path: &Path) -> bool {
        self.directory.exists(path)
    }

    fn acquire_lock(&self, lock: &Lock) -> result::Result<DirectoryLock, LockError> {
        self.directory.acquire_lock(lock)
    }

    fn watch(&self, watch_callback: WatchCallback) -> WatchHandle {
        self.directory.watch(watch_callback)
    }
}

impl Clone for ManagedDirectory {
    fn clone(&self) -> ManagedDirectory {
        ManagedDirectory {
            directory: self.directory.box_clone(),
            meta_informations: Arc::clone(&self.meta_informations),
        }
    }
}

#[cfg(test)]
mod tests {

    #[cfg(feature = "mmap")]
    mod mmap_specific {

        use super::super::*;
        use once_cell::sync::Lazy;
        use std::path::Path;
        use tempdir::TempDir;

        static TEST_PATH1: Lazy<&'static Path> = Lazy::new(|| Path::new("some_path_for_test"));
        static TEST_PATH2: Lazy<&'static Path> = Lazy::new(|| Path::new("some_path_for_test2"));

        use crate::directory::MmapDirectory;
        use std::io::Write;

        #[test]
        fn test_managed_directory() {
            let tempdir = TempDir::new("index").unwrap();
            let tempdir_path = PathBuf::from(tempdir.path());
            {
                let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
                let mut managed_directory = ManagedDirectory::wrap(mmap_directory).unwrap();
                {
                    let mut write_file = managed_directory.open_write(*TEST_PATH1).unwrap();
                    write_file.flush().unwrap();
                }
                {
                    managed_directory
                        .atomic_write(*TEST_PATH2, &vec![0u8, 1u8])
                        .unwrap();
                }
                {
                    assert!(managed_directory.exists(*TEST_PATH1));
                    assert!(managed_directory.exists(*TEST_PATH2));
                }
                {
                    let living_files: HashSet<PathBuf> =
                        [TEST_PATH1.to_owned()].into_iter().cloned().collect();
                    managed_directory.garbage_collect(|| living_files);
                }
                {
                    assert!(managed_directory.exists(*TEST_PATH1));
                    assert!(!managed_directory.exists(*TEST_PATH2));
                }
            }
            {
                let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
                let mut managed_directory = ManagedDirectory::wrap(mmap_directory).unwrap();
                {
                    assert!(managed_directory.exists(*TEST_PATH1));
                    assert!(!managed_directory.exists(*TEST_PATH2));
                }
                {
                    let living_files: HashSet<PathBuf> = HashSet::new();
                    managed_directory.garbage_collect(|| living_files);
                }
                {
                    assert!(!managed_directory.exists(*TEST_PATH1));
                    assert!(!managed_directory.exists(*TEST_PATH2));
                }
            }
        }

        #[test]
        fn test_managed_directory_gc_while_mmapped() {
            let tempdir = TempDir::new("index").unwrap();
            let tempdir_path = PathBuf::from(tempdir.path());
            let living_files = HashSet::new();

            let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
            let mut managed_directory = ManagedDirectory::wrap(mmap_directory).unwrap();
            managed_directory
                .atomic_write(*TEST_PATH1, &vec![0u8, 1u8])
                .unwrap();
            assert!(managed_directory.exists(*TEST_PATH1));

            let _mmap_read = managed_directory.open_read(*TEST_PATH1).unwrap();
            managed_directory.garbage_collect(|| living_files.clone());
            if cfg!(target_os = "windows") {
                // On Windows, gc should try and fail the file as it is mmapped.
                assert!(managed_directory.exists(*TEST_PATH1));
                // unmap should happen here.
                drop(_mmap_read);
                // The file should still be in the list of managed file and
                // eventually be deleted once mmap is released.
                managed_directory.garbage_collect(|| living_files);
                assert!(!managed_directory.exists(*TEST_PATH1));
            } else {
                assert!(!managed_directory.exists(*TEST_PATH1));
            }
        }

    }

}
<commit_msg>Small code cleaning<commit_after>use crate::core::MANAGED_FILEPATH;
use crate::directory::error::{DeleteError, IOError, LockError, OpenReadError, OpenWriteError};
use crate::directory::DirectoryLock;
use crate::directory::Lock;
use crate::directory::META_LOCK;
use crate::directory::{ReadOnlySource, WritePtr};
use crate::directory::{WatchCallback, WatchHandle};
use crate::error::DataCorruption;
use crate::Directory;
use crate::Result;
use serde_json;
use std::collections::HashSet;
use std::io;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::result;
use std::sync::RwLockWriteGuard;
use std::sync::{Arc, RwLock};

/// Returns true iff the file is "managed".
/// Non-managed file are not subject to garbage collection.
///
/// Filenames that starts by a "." -typically locks-
/// are not managed.
fn is_managed(path: &Path) -> bool {
    path.to_str()
        .map(|p_str| !p_str.starts_with('.'))
        .unwrap_or(true)
}

/// Wrapper of directories that keeps track of files created by Tantivy.
///
/// A managed directory is just a wrapper of a directory
/// that keeps a (persisted) list of the files that
/// have been created (and not deleted) by tantivy so far.
///
/// Thanks to this list, it implements a `garbage_collect` method
/// that removes the files that were created by tantivy and are not
/// useful anymore.
#[derive(Debug)]
pub struct ManagedDirectory {
    directory: Box<dyn Directory>,
    meta_informations: Arc<RwLock<MetaInformation>>,
}

#[derive(Debug, Default)]
struct MetaInformation {
    managed_paths: HashSet<PathBuf>,
}

/// Saves the file containing the list of existing files
/// that were created by tantivy.
fn save_managed_paths(
    directory: &mut dyn Directory,
    wlock: &RwLockWriteGuard<'_, MetaInformation>,
) -> io::Result<()> {
    let mut w = serde_json::to_vec(&wlock.managed_paths)?;
    writeln!(&mut w)?;
    directory.atomic_write(&MANAGED_FILEPATH, &w[..])?;
    Ok(())
}

impl ManagedDirectory {
    /// Wraps a directory as managed directory.
    pub fn wrap<Dir: Directory>(directory: Dir) -> Result<ManagedDirectory> {
        match directory.atomic_read(&MANAGED_FILEPATH) {
            Ok(data) => {
                let managed_files_json = String::from_utf8_lossy(&data);
                let managed_files: HashSet<PathBuf> = serde_json::from_str(&managed_files_json)
                    .map_err(|e| {
                        DataCorruption::new(
                            MANAGED_FILEPATH.to_path_buf(),
                            format!("Managed file cannot be deserialized: {:?}. ", e),
                        )
                    })?;
                Ok(ManagedDirectory {
                    directory: Box::new(directory),
                    meta_informations: Arc::new(RwLock::new(MetaInformation {
                        managed_paths: managed_files,
                    })),
                })
            }
            Err(OpenReadError::FileDoesNotExist(_)) => Ok(ManagedDirectory {
                directory: Box::new(directory),
                meta_informations: Arc::default(),
            }),
            Err(OpenReadError::IOError(e)) => Err(From::from(e)),
        }
    }

    /// Garbage collect unused files.
    ///
    /// Removes the files that were created by `tantivy` and are not
    /// used by any segment anymore.
    ///
    /// * `living_files` - List of files that are still used by the index.
    ///
    /// The use a callback ensures that the list of living_files is computed
    /// while we hold the lock on meta.
    ///
    /// This method does not panick nor returns errors.
    /// If a file cannot be deleted (for permission reasons for instance)
    /// an error is simply logged, and the file remains in the list of managed
    /// files.
    pub fn garbage_collect<L: FnOnce() -> HashSet<PathBuf>>(&mut self, get_living_files: L) {
        info!("Garbage collect");
        let mut files_to_delete = vec![];

        // It is crucial to get the living files after acquiring the
        // read lock of meta informations. That way, we
        // avoid the following scenario.
        //
        // 1) we get the list of living files.
        // 2) someone creates a new file.
        // 3) we start garbage collection and remove this file
        // even though it is a living file.
        //
        // releasing the lock as .delete() will use it too.
        {
            let meta_informations_rlock = self
                .meta_informations
                .read()
                .expect("Managed directory rlock poisoned in garbage collect.");

            // The point of this second "file" lock is to enforce the following scenario
            // 1) process B tries to load a new set of searcher.
            // The list of segments is loaded
            // 2) writer change meta.json (for instance after a merge or a commit)
            // 3) gc kicks in.
            // 4) gc removes a file that was useful for process B, before process B opened it.
            if let Ok(_meta_lock) = self.acquire_lock(&META_LOCK) {
                let living_files = get_living_files();
                for managed_path in &meta_informations_rlock.managed_paths {
                    if !living_files.contains(managed_path) {
                        files_to_delete.push(managed_path.clone());
                    }
                }
            }
        }

        let mut deleted_files = vec![];
        {
            for file_to_delete in files_to_delete {
                match self.delete(&file_to_delete) {
                    Ok(_) => {
                        info!("Deleted {:?}", file_to_delete);
                        deleted_files.push(file_to_delete);
                    }
                    Err(file_error) => {
                        match file_error {
                            DeleteError::FileDoesNotExist(_) => {
                                deleted_files.push(file_to_delete);
                            }
                            DeleteError::IOError(_) => {
                                if !cfg!(target_os = "windows") {
                                    // On windows, delete is expected to fail if the file
                                    // is mmapped.
                                    error!("Failed to delete {:?}", file_to_delete);
                                }
                            }
                        }
                    }
                }
            }
        }

        if !deleted_files.is_empty() {
            // update the list of managed files by removing
            // the file that were removed.
            let mut meta_informations_wlock = self
                .meta_informations
                .write()
                .expect("Managed directory wlock poisoned (2).");
            {
                let managed_paths_write = &mut meta_informations_wlock.managed_paths;
                for delete_file in &deleted_files {
                    managed_paths_write.remove(delete_file);
                }
            }
            if save_managed_paths(self.directory.as_mut(), &meta_informations_wlock).is_err() {
                error!("Failed to save the list of managed files.");
            }
        }
    }

    /// Registers a file as managed
    ///
    /// This method must be called before the file is
    /// actually created to ensure that a failure between
    /// registering the filepath and creating the file
    /// will not lead to garbage files that will
    /// never get removed.
    ///
    /// File starting by "." are reserved to locks.
    /// They are not managed and cannot be subjected
    /// to garbage collection.
    fn register_file_as_managed(&mut self, filepath: &Path) -> io::Result<()> {
        // Files starting by "." (e.g. lock files) are not managed.
        if !is_managed(filepath) {
            return Ok(());
        }
        let mut meta_wlock = self
            .meta_informations
            .write()
            .expect("Managed file lock poisoned");
        let has_changed = meta_wlock.managed_paths.insert(filepath.to_owned());
        if has_changed {
            save_managed_paths(self.directory.as_mut(), &meta_wlock)?;
        }
        Ok(())
    }
}

impl Directory for ManagedDirectory {
    fn open_read(&self, path: &Path) -> result::Result<ReadOnlySource, OpenReadError> {
        self.directory.open_read(path)
    }

    fn open_write(&mut self, path: &Path) -> result::Result<WritePtr, OpenWriteError> {
        self.register_file_as_managed(path)
            .map_err(|e| IOError::with_path(path.to_owned(), e))?;
        self.directory.open_write(path)
    }

    fn atomic_write(&mut self, path: &Path, data: &[u8]) -> io::Result<()> {
        self.register_file_as_managed(path)?;
        self.directory.atomic_write(path, data)
    }

    fn atomic_read(&self, path: &Path) -> result::Result<Vec<u8>, OpenReadError> {
        self.directory.atomic_read(path)
    }

    fn delete(&self, path: &Path) -> result::Result<(), DeleteError> {
        self.directory.delete(path)
    }

    fn exists(&self, path: &Path) -> bool {
        self.directory.exists(path)
    }

    fn acquire_lock(&self, lock: &Lock) -> result::Result<DirectoryLock, LockError> {
        self.directory.acquire_lock(lock)
    }

    fn watch(&self, watch_callback: WatchCallback) -> WatchHandle {
        self.directory.watch(watch_callback)
    }
}

impl Clone for ManagedDirectory {
    fn clone(&self) -> ManagedDirectory {
        ManagedDirectory {
            directory: self.directory.box_clone(),
            meta_informations: Arc::clone(&self.meta_informations),
        }
    }
}

#[cfg(feature = "mmap")]
#[cfg(test)]
mod tests_mmap_specific {

    use crate::directory::{Directory, ManagedDirectory, MmapDirectory};
    use std::collections::HashSet;
    use std::io::Write;
    use std::path::{Path, PathBuf};
    use tempdir::TempDir;

    #[test]
    fn test_managed_directory() {
        let tempdir = TempDir::new("tantivy-test").unwrap();
        let tempdir_path = PathBuf::from(tempdir.path());

        let test_path1: &'static Path = Path::new("some_path_for_test");
        let test_path2: &'static Path = Path::new("some_path_for_test_2");
        {
            let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
            let mut managed_directory = ManagedDirectory::wrap(mmap_directory).unwrap();
            let mut write_file = managed_directory.open_write(test_path1).unwrap();
            write_file.flush().unwrap();
            managed_directory
                .atomic_write(test_path2, &[0u8, 1u8])
                .unwrap();
            assert!(managed_directory.exists(test_path1));
            assert!(managed_directory.exists(test_path2));
            let living_files: HashSet<PathBuf> =
                [test_path1.to_owned()].into_iter().cloned().collect();
            managed_directory.garbage_collect(|| living_files);
            assert!(managed_directory.exists(test_path1));
            assert!(!managed_directory.exists(test_path2));
        }
        {
            let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
            let mut managed_directory = ManagedDirectory::wrap(mmap_directory).unwrap();
            assert!(managed_directory.exists(test_path1));
            assert!(!managed_directory.exists(test_path2));
            let living_files: HashSet<PathBuf> = HashSet::new();
            managed_directory.garbage_collect(|| living_files);
            assert!(!managed_directory.exists(test_path1));
            assert!(!managed_directory.exists(test_path2));
        }
    }

    #[test]
    fn test_managed_directory_gc_while_mmapped() {
        let test_path1: &'static Path = Path::new("some_path_for_test");

        let tempdir = TempDir::new("index").unwrap();
        let tempdir_path = PathBuf::from(tempdir.path());
        let living_files = HashSet::new();

        let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
        let mut managed_directory = ManagedDirectory::wrap(mmap_directory).unwrap();
        managed_directory
            .atomic_write(test_path1, &vec![0u8, 1u8])
            .unwrap();
        assert!(managed_directory.exists(test_path1));

        let _mmap_read = managed_directory.open_read(test_path1).unwrap();
        managed_directory.garbage_collect(|| living_files.clone());
        if cfg!(target_os = "windows") {
            // On Windows, gc should try and fail the file as it is mmapped.
            assert!(managed_directory.exists(test_path1));
            // unmap should happen here.
            drop(_mmap_read);
            // The file should still be in the list of managed file and
            // eventually be deleted once mmap is released.
            managed_directory.garbage_collect(|| living_files);
            assert!(!managed_directory.exists(test_path1));
        } else {
            assert!(!managed_directory.exists(test_path1));
        }
    }

}
<|endoftext|>