<commit_before>use std::collections::HashMap;
use std::fmt;

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

impl Metadata {
  pub fn new(is_last: bool, length: u32, data: Data) -> Self {
    Metadata {
      is_last: is_last,
      length: length,
      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],
    }
  }

  #[inline]
  pub fn is_varied_block_size(&self) -> bool {
    self.min_block_size != self.max_block_size
  }

  #[inline]
  pub fn is_fixed_block_size(&self) -> bool {
    self.min_block_size == self.max_block_size
  }
}

/// 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,
}

impl fmt::Display for PictureType {
  fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
    write!(formatter, "{}", match *self {
      PictureType::Other              => "Other",
      PictureType::FileIconStandard   => "File Icon (standard)",
      PictureType::FileIcon           => "File Icon",
      PictureType::FrontCover         => "Cover (front)",
      PictureType::BackCover          => "Cover (back)",
      PictureType::LeafletPage        => "Leaflet Page",
      PictureType::Media              => "Media",
      PictureType::LeadArtist         => "Lead Artist",
      PictureType::Artist             => "Arist",
      PictureType::Conductor          => "Conductor",
      PictureType::Band               => "Band",
      PictureType::Composer           => "Composer",
      PictureType::Lyricist           => "Lyricist",
      PictureType::RecordingLocation  => "Recoding Location",
      PictureType::DuringRecording    => "During Recording",
      PictureType::DuringPerformace   => "During Performace",
      PictureType::VideoScreenCapture => "Video Screen Capture",
      PictureType::Fish               => "Fish",
      PictureType::Illustration       => "Illustration",
      PictureType::BandLogo           => "Band Logo",
      PictureType::PublisherLogo      => "Publisher Logo",
    })
  }
}

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

  #[test]
  fn test_is_varied_block_size() {
    let mut info = StreamInfo::new();

    info.min_block_size = 512;
    info.max_block_size = 1024;

    assert!(info.is_varied_block_size());

    info.min_block_size = 2048;
    info.max_block_size = 2048;

    assert!(!info.is_varied_block_size());
  }

  #[test]
  fn test_is_fixed_block_size() {
    let mut info = StreamInfo::new();

    info.min_block_size = 512;
    info.max_block_size = 512;

    assert!(info.is_fixed_block_size());

    info.min_block_size = 1024;
    info.max_block_size = 2048;

    assert!(!info.is_fixed_block_size());
  }
}
<commit_msg>Add `Metadata::is_last` method<commit_after>use std::collections::HashMap;
use std::fmt;

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

impl Metadata {
  pub fn new(is_last: bool, length: u32, data: Data) -> Self {
    Metadata {
      is_last: is_last,
      length: length,
      data: data,
    }
  }

  /// Return whether the current metadata block is the last.
  #[inline]
  pub fn is_last(&self) -> bool {
    self.is_last
  }
}

/// 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],
    }
  }

  #[inline]
  pub fn is_varied_block_size(&self) -> bool {
    self.min_block_size != self.max_block_size
  }

  #[inline]
  pub fn is_fixed_block_size(&self) -> bool {
    self.min_block_size == self.max_block_size
  }
}

/// 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,
}

impl fmt::Display for PictureType {
  fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
    write!(formatter, "{}", match *self {
      PictureType::Other              => "Other",
      PictureType::FileIconStandard   => "File Icon (standard)",
      PictureType::FileIcon           => "File Icon",
      PictureType::FrontCover         => "Cover (front)",
      PictureType::BackCover          => "Cover (back)",
      PictureType::LeafletPage        => "Leaflet Page",
      PictureType::Media              => "Media",
      PictureType::LeadArtist         => "Lead Artist",
      PictureType::Artist             => "Arist",
      PictureType::Conductor          => "Conductor",
      PictureType::Band               => "Band",
      PictureType::Composer           => "Composer",
      PictureType::Lyricist           => "Lyricist",
      PictureType::RecordingLocation  => "Recoding Location",
      PictureType::DuringRecording    => "During Recording",
      PictureType::DuringPerformace   => "During Performace",
      PictureType::VideoScreenCapture => "Video Screen Capture",
      PictureType::Fish               => "Fish",
      PictureType::Illustration       => "Illustration",
      PictureType::BandLogo           => "Band Logo",
      PictureType::PublisherLogo      => "Publisher Logo",
    })
  }
}

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

  #[test]
  fn test_is_varied_block_size() {
    let mut info = StreamInfo::new();

    info.min_block_size = 512;
    info.max_block_size = 1024;

    assert!(info.is_varied_block_size());

    info.min_block_size = 2048;
    info.max_block_size = 2048;

    assert!(!info.is_varied_block_size());
  }

  #[test]
  fn test_is_fixed_block_size() {
    let mut info = StreamInfo::new();

    info.min_block_size = 512;
    info.max_block_size = 512;

    assert!(info.is_fixed_block_size());

    info.min_block_size = 1024;
    info.max_block_size = 2048;

    assert!(!info.is_fixed_block_size());
  }
}
<|endoftext|>


<commit_before>extern crate libc;
use self::libc::{c_double, c_int, size_t};

use super::{Function, Parameterised, Linear, VFunction, QFunction};

use ndarray::{Axis, Array1};
use geometry::{Space, RegularSpace};
use geometry::dimensions::Continuous;

#[link(name="tiles", kind="static")]
extern {
    fn tiles(tile_indices: *mut size_t, nt: c_int, memory: c_int,
             floats: *const c_double, nf: c_int,
             ints: *const c_int, ni: c_int);
}

pub struct SuttonTiles {
    weights: Array1<c_double>,

    n_outputs: usize,
    n_tilings: i32,
    memory_size: i32,

    int_array: [c_int; 1]
}

impl SuttonTiles {
    pub fn new(n_tilings: i32, memory_size: i32, n_outputs: usize) -> Self {
        SuttonTiles {
            weights: Array1::<c_double>::zeros(memory_size as usize),

            n_outputs: n_outputs,
            n_tilings: n_tilings,
            memory_size: memory_size,

            int_array: [0],
        }
    }

    fn load_tiles(&self, floats: &[c_double], ints: &[c_int]) -> Vec<size_t> {
        let mut ti = vec![0; self.n_tilings as usize];

        unsafe {
            tiles(ti.as_mut_ptr(), self.n_tilings, self.memory_size,
                  floats.as_ptr(), floats.len() as c_int,
                  ints.as_ptr(), ints.len() as c_int);
        }

        ti
    }

    fn evaluate_index(&self, input: &Vec<f64>, index: c_int) -> f64 {
        let tiles = self.load_tiles(input, &[index]);

        self.weights.select(Axis(0), tiles.as_slice()).scalar_sum() / self.n_tilings as f64
    }
}

impl Function<Vec<f64>, f64> for SuttonTiles {
    fn evaluate(&self, input: &Vec<f64>) -> f64 {
        self.evaluate_index(input, 0)
    }
}

impl Function<Vec<f64>, Vec<f64>> for SuttonTiles {
    fn evaluate(&self, input: &Vec<f64>) -> Vec<f64> {
        (0..self.n_outputs).map(|c| self.evaluate_index(input, c as c_int)).collect()
    }
}

impl Parameterised<Vec<f64>, f64> for SuttonTiles {
    fn update(&mut self, input: &Vec<f64>, error: f64) {
        self.int_array[0] = 0;

        for t in self.load_tiles(input, &self.int_array) {
            self.weights[t] += error;
        }
    }
}

impl Parameterised<Vec<f64>, Vec<f64>> for SuttonTiles {
    fn update(&mut self, input: &Vec<f64>, errors: Vec<f64>) {
        for c in 0..self.n_outputs {
            <Self as QFunction<RegularSpace<Continuous>>>::update_action(self, input, c, errors[c]);
        }
    }
}

// TODO: Implement Linear - problem is that phi will be a function of the state
//       and action for this implementation of tile coding.

impl VFunction<RegularSpace<Continuous>> for SuttonTiles {}

impl QFunction<RegularSpace<Continuous>> for SuttonTiles
{
    fn evaluate_action(&self, input: &Vec<f64>, action: usize) -> f64 {
        self.evaluate_index(input, action as c_int)
    }

    fn update_action(&mut self, input: &Vec<f64>, action: usize, error: f64) {
        self.int_array[0] = action as c_int;

        for t in self.load_tiles(input, &self.int_array) {
            self.weights[t] += error;
        }
    }
}

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

    use fa::{Function, Parameterised};
    use geometry::RegularSpace;
    use geometry::dimensions::Partition;

    #[test]
    fn test_simple() {
        let mut t = SuttonTiles::new(1, 1000, 1);

        t.update(&vec![1.5], 25.5);

        let out: f64 = t.evaluate(&vec![1.5]);
        assert_eq!(out, 25.5);

        t.update(&vec![1.5], -12.75);

        let out: f64 = t.evaluate(&vec![1.5]);
        assert_eq!(out, 12.75);
    }
}
<commit_msg>Added test for unit generalisation in SuttonTiles<commit_after>extern crate libc;
use self::libc::{c_double, c_int, size_t};

use super::{Function, Parameterised, Linear, VFunction, QFunction};

use ndarray::{Axis, Array1};
use geometry::{Space, RegularSpace};
use geometry::dimensions::Continuous;

#[link(name="tiles", kind="static")]
extern {
    fn tiles(tile_indices: *mut size_t, nt: c_int, memory: c_int,
             floats: *const c_double, nf: c_int,
             ints: *const c_int, ni: c_int);
}

pub struct SuttonTiles {
    weights: Array1<c_double>,

    n_outputs: usize,
    n_tilings: i32,
    memory_size: i32,

    int_array: [c_int; 1]
}

impl SuttonTiles {
    pub fn new(n_tilings: i32, memory_size: i32, n_outputs: usize) -> Self {
        SuttonTiles {
            weights: Array1::<c_double>::zeros(memory_size as usize),

            n_outputs: n_outputs,
            n_tilings: n_tilings,
            memory_size: memory_size,

            int_array: [0],
        }
    }

    fn load_tiles(&self, floats: &[c_double], ints: &[c_int]) -> Vec<size_t> {
        let mut ti = vec![0; self.n_tilings as usize];

        unsafe {
            tiles(ti.as_mut_ptr(), self.n_tilings, self.memory_size,
                  floats.as_ptr(), floats.len() as c_int,
                  ints.as_ptr(), ints.len() as c_int);
        }

        ti
    }

    fn evaluate_index(&self, input: &Vec<f64>, index: c_int) -> f64 {
        let tiles = self.load_tiles(input, &[index]);

        self.weights.select(Axis(0), tiles.as_slice()).scalar_sum() / self.n_tilings as f64
    }
}

impl Function<Vec<f64>, f64> for SuttonTiles {
    fn evaluate(&self, input: &Vec<f64>) -> f64 {
        self.evaluate_index(input, 0)
    }
}

impl Function<Vec<f64>, Vec<f64>> for SuttonTiles {
    fn evaluate(&self, input: &Vec<f64>) -> Vec<f64> {
        (0..self.n_outputs).map(|c| self.evaluate_index(input, c as c_int)).collect()
    }
}

impl Parameterised<Vec<f64>, f64> for SuttonTiles {
    fn update(&mut self, input: &Vec<f64>, error: f64) {
        self.int_array[0] = 0;

        for t in self.load_tiles(input, &self.int_array) {
            self.weights[t] += error;
        }
    }
}

impl Parameterised<Vec<f64>, Vec<f64>> for SuttonTiles {
    fn update(&mut self, input: &Vec<f64>, errors: Vec<f64>) {
        for c in 0..self.n_outputs {
            <Self as QFunction<RegularSpace<Continuous>>>::update_action(self, input, c, errors[c]);
        }
    }
}

// TODO: Implement Linear - problem is that phi will be a function of the state
//       and action for this implementation of tile coding.

impl VFunction<RegularSpace<Continuous>> for SuttonTiles {}

impl QFunction<RegularSpace<Continuous>> for SuttonTiles
{
    fn evaluate_action(&self, input: &Vec<f64>, action: usize) -> f64 {
        self.evaluate_index(input, action as c_int)
    }

    fn update_action(&mut self, input: &Vec<f64>, action: usize, error: f64) {
        self.int_array[0] = action as c_int;

        for t in self.load_tiles(input, &self.int_array) {
            self.weights[t] += error;
        }
    }
}

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

    use fa::{Function, Parameterised};
    use geometry::RegularSpace;
    use geometry::dimensions::Partition;

    #[test]
    fn test_simple() {
        let mut t = SuttonTiles::new(1, 1000, 1);

        t.update(&vec![1.5], 25.5);

        let out: f64 = t.evaluate(&vec![1.5]);
        assert_eq!(out, 25.5);

        t.update(&vec![1.5], -12.75);

        let out: f64 = t.evaluate(&vec![1.5]);
        assert_eq!(out, 12.75);
    }

    #[test]
    fn test_generalisation() {
        let mut t = SuttonTiles::new(1, 1000, 1);

        t.update(&vec![0.5], 1.0);

        for i in 1..10 {
            let out: f64 = t.evaluate(&vec![i as f64 / 10.0]);
            assert_eq!(out, 1.0);
        }

        let out: f64 = t.evaluate(&vec![-0.000001]);
        assert_eq!(out, 0.0);

        let out: f64 = t.evaluate(&vec![1.000001]);
        assert_eq!(out, 0.0);
    }
}
<|endoftext|>


<commit_before>use std::cmp;
use decoders::RawImage;
use imageops::{OpBuffer,ImageOp,Pipeline};

#[derive(Clone, Debug)]
pub struct OpBaseCurve {
  xs: Vec<f32>,
  ys: Vec<f32>,
}

impl OpBaseCurve {
  pub fn new(_img: &RawImage) -> OpBaseCurve {
    OpBaseCurve{
      xs: vec![0.0, 0.30, 0.5, 0.70, 1.0],
      ys: vec![0.0, 0.25, 0.5, 0.75, 1.0],
    }
  }
}

impl ImageOp for OpBaseCurve {
  fn name(&self) -> &str {"basecurve"}
  fn run(&self, _pipeline: &Pipeline, buf: &OpBuffer) -> OpBuffer {
    let mut buf = buf.clone();
    let func = SplineFunc::new(&self.xs, &self.ys);

    buf.mutate_lines(&(|line: &mut [f32], _| {
      for pix in line.chunks_mut(3) {
        pix[0] = func.interpolate(pix[0]);
        pix[1] = pix[1];
        pix[2] = pix[2];
      }
    }));

    buf
  }
}

struct SplineFunc<'a> {
  xs: &'a [f32],
  ys: &'a [f32],
  c1s: Vec<f32>,
  c2s: Vec<f32>,
  c3s: Vec<f32>,
}

impl<'a> SplineFunc<'a> {
  // Monotone cubic interpolation code adapted from the Javascript example in Wikipedia
  fn new(xs: &'a[f32], ys: &'a[f32]) -> SplineFunc<'a> {
    if xs.len() != ys.len() { panic!("Different number of Xs and Ys for Spline"); }
    if xs.len() < 2 { panic!("Need at least 2 points for Spline"); }

	  // Get consecutive differences and slopes
    let mut dxs = Vec::new();
    let mut dys = Vec::new();
    let mut slopes = Vec::new();
    for i in 0..(xs.len()-1) {
      let dx = xs[i+1] - xs[i];
      let dy = ys[i+1] - ys[i];
      dxs.push(dx);
      dys.push(dy);
      slopes.push(dy/dx);
    }

	  // Get degree-1 coefficients
    let mut c1s = vec![slopes[0]];
    for i in 0..(dxs.len()-1) {
      let m = slopes[i];
      let next = slopes[i+1];
      if m*next <= 0.0 {
        c1s.push(0.0);
      } else {
        let dx = dxs[i];
        let dxnext = dxs[i+1];
        let common = dx + dxnext;
        c1s.push(3.0*common/((common+dxnext)/m + (common + dx)/next));
      }
    }
    c1s.push(slopes[slopes.len()-1]);

	  // Get degree-2 and degree-3 coefficients
    let mut c2s = Vec::new();
    let mut c3s = Vec::new();
    for i in 0..(c1s.len()-1) {
      let c1 = c1s[i];
      let slope = slopes[i];
      let invdx = 1.0 / dxs[i];
      let common = c1+c1s[i+1] - slope - slope;
      c2s.push((slope-c1-common)*invdx);
      c3s.push(common*invdx*invdx);
    }

    SplineFunc {
      xs: xs,
      ys: ys,
        c1s: c1s,
      c2s: c2s,
      c3s: c3s,
    }
  }

  fn interpolate(&self, val: f32) -> f32 {
		// Anything at or beyond the last value returns the last value
    let end = self.xs[self.xs.len()-1];
    if val >= end {
      return self.ys[self.ys.len()-1];
    }

		// Search for the interval x is in, returning the corresponding y if x is one of the original xs
    let mut low: isize = 0;
    let mut mid: isize;
    let mut high: isize = (self.c3s.len() - 1) as isize;

		while low <= high {
      mid = (low+high)/2;
      let xhere = self.xs[mid as usize];
      if xhere < val { low = mid + 1; }
      else if xhere > val { high = mid - 1; }
      else { return self.ys[mid as usize] }
		}
    let i = cmp::max(0, high) as usize;

		// Interpolate
    let diff = val - self.xs[i];

    self.ys[i] + self.c1s[i]*diff + self.c2s[i]*diff*diff + self.c3s[i]*diff*diff*diff
	}
}
<commit_msg>curves: get rid of tabs<commit_after>use std::cmp;
use decoders::RawImage;
use imageops::{OpBuffer,ImageOp,Pipeline};

#[derive(Clone, Debug)]
pub struct OpBaseCurve {
  xs: Vec<f32>,
  ys: Vec<f32>,
}

impl OpBaseCurve {
  pub fn new(_img: &RawImage) -> OpBaseCurve {
    OpBaseCurve{
      xs: vec![0.0, 0.30, 0.5, 0.70, 1.0],
      ys: vec![0.0, 0.25, 0.5, 0.75, 1.0],
    }
  }
}

impl ImageOp for OpBaseCurve {
  fn name(&self) -> &str {"basecurve"}
  fn run(&self, _pipeline: &Pipeline, buf: &OpBuffer) -> OpBuffer {
    let mut buf = buf.clone();
    let func = SplineFunc::new(&self.xs, &self.ys);

    buf.mutate_lines(&(|line: &mut [f32], _| {
      for pix in line.chunks_mut(3) {
        pix[0] = func.interpolate(pix[0]);
        pix[1] = pix[1];
        pix[2] = pix[2];
      }
    }));

    buf
  }
}

struct SplineFunc<'a> {
  xs: &'a [f32],
  ys: &'a [f32],
  c1s: Vec<f32>,
  c2s: Vec<f32>,
  c3s: Vec<f32>,
}

impl<'a> SplineFunc<'a> {
  // Monotone cubic interpolation code adapted from the Javascript example in Wikipedia
  fn new(xs: &'a[f32], ys: &'a[f32]) -> SplineFunc<'a> {
    if xs.len() != ys.len() { panic!("Different number of Xs and Ys for Spline"); }
    if xs.len() < 2 { panic!("Need at least 2 points for Spline"); }

    // Get consecutive differences and slopes
    let mut dxs = Vec::new();
    let mut dys = Vec::new();
    let mut slopes = Vec::new();
    for i in 0..(xs.len()-1) {
      let dx = xs[i+1] - xs[i];
      let dy = ys[i+1] - ys[i];
      dxs.push(dx);
      dys.push(dy);
      slopes.push(dy/dx);
    }

    // Get degree-1 coefficients
    let mut c1s = vec![slopes[0]];
    for i in 0..(dxs.len()-1) {
      let m = slopes[i];
      let next = slopes[i+1];
      if m*next <= 0.0 {
        c1s.push(0.0);
      } else {
        let dx = dxs[i];
        let dxnext = dxs[i+1];
        let common = dx + dxnext;
        c1s.push(3.0*common/((common+dxnext)/m + (common + dx)/next));
      }
    }
    c1s.push(slopes[slopes.len()-1]);

    // Get degree-2 and degree-3 coefficients
    let mut c2s = Vec::new();
    let mut c3s = Vec::new();
    for i in 0..(c1s.len()-1) {
      let c1 = c1s[i];
      let slope = slopes[i];
      let invdx = 1.0 / dxs[i];
      let common = c1+c1s[i+1] - slope - slope;
      c2s.push((slope-c1-common)*invdx);
      c3s.push(common*invdx*invdx);
    }

    SplineFunc {
      xs: xs,
      ys: ys,
        c1s: c1s,
      c2s: c2s,
      c3s: c3s,
    }
  }

  fn interpolate(&self, val: f32) -> f32 {
    // Anything at or beyond the last value returns the last value
    let end = self.xs[self.xs.len()-1];
    if val >= end {
      return self.ys[self.ys.len()-1];
    }

    // Search for the interval x is in, returning the corresponding y if x is one of the original xs
    let mut low: isize = 0;
    let mut mid: isize;
    let mut high: isize = (self.c3s.len() - 1) as isize;

    while low <= high {
      mid = (low+high)/2;
      let xhere = self.xs[mid as usize];
      if xhere < val { low = mid + 1; }
      else if xhere > val { high = mid - 1; }
      else { return self.ys[mid as usize] }
    }
    let i = cmp::max(0, high) as usize;

    // Interpolate
    let diff = val - self.xs[i];

    self.ys[i] + self.c1s[i]*diff + self.c2s[i]*diff*diff + self.c3s[i]*diff*diff*diff
  }
}
<|endoftext|>


<commit_before>//! WebAssembly interpreter module.

/// Custom user error.
pub trait UserError: 'static + ::std::fmt::Display + ::std::fmt::Debug + Clone + PartialEq {
}

/// Internal interpreter error.
#[derive(Debug, Clone, PartialEq)]
pub enum Error<E> where E: UserError {
	/// Program-level error.
	Program(String),
	/// Validation error.
	Validation(String),
	/// Initialization error.
	Initialization(String),
	/// Function-level error.
	Function(String),
	/// Table-level error.
	Table(String),
	/// Memory-level error.
	Memory(String),
	/// Variable-level error.
	Variable(String),
	/// Global-level error.
	Global(String),
	/// Local-level error.
	Local(String),
	/// Stack-level error.
	Stack(String),
	/// Value-level error.
	Value(String),
	/// Interpreter (code) error.
	Interpreter(String),
	/// Env module error.
	Env(String),
	/// Native module error.
	Native(String),
	/// Trap.
	Trap(String),
	/// Custom user error.
	User(E),
}

impl<E> Into<String> for Error<E> where E: UserError {
	fn into(self) -> String {
		match self {
			Error::Program(s) => s,
			Error::Validation(s) => s,
			Error::Initialization(s) => s,
			Error::Function(s) => s,
			Error::Table(s) => s,
			Error::Memory(s) => s,
			Error::Variable(s) => s,
			Error::Global(s) => s,
			Error::Local(s) => s,
			Error::Stack(s) => s,
			Error::Interpreter(s) => s,
			Error::Value(s) => s,
			Error::Env(s) => s,
			Error::Native(s) => s,
			Error::Trap(s) => format!("trap: {}", s),
			Error::User(e) => format!("user: {}", e),
		}
	}
}

/// Dummy user error.
#[derive(Debug, Clone, PartialEq)]
pub struct DummyUserError;

impl UserError for DummyUserError {}

impl ::std::fmt::Display for DummyUserError {
	fn fmt(&self, _f: &mut ::std::fmt::Formatter) -> Result<(), ::std::fmt::Error> { Ok(()) }
}

mod env;
mod env_native;
mod imports;
mod memory;
mod module;
mod program;
mod runner;
mod stack;
mod table;
mod validator;
mod value;
mod variable;

#[cfg(test)]
mod tests;

pub use self::memory::MemoryInstance;
pub use self::module::{ModuleInstance, ModuleInstanceInterface,
	ItemIndex, ExportEntryType, CallerContext, ExecutionParams, FunctionSignature};
pub use self::table::TableInstance;
pub use self::program::ProgramInstance;
pub use self::value::RuntimeValue;
pub use self::variable::{VariableInstance, VariableType, ExternalVariableValue};
pub use self::env_native::{env_native_module, UserDefinedElements, UserFunctionExecutor, UserFunctionDescriptor};
pub use self::env::EnvParams;

/// Default type of Error if you do not need any custom user errors.
pub type DummyError = Error<DummyUserError>;

/// Default type of ProgramInstance if you do not need any custom user errors.
/// To work with custom user errors or interpreter internals, use CustomProgramInstance.
pub type DefaultProgramInstance = self::program::ProgramInstance<DummyUserError>;

/// Default type of ModuleInstance if you do not need any custom user errors.
/// To work with custom user errors or interpreter internals, use CustomModuleInstance.
pub type DefaultModuleInstance = self::module::ModuleInstance<DummyUserError>;

/// Default type of ModuleInstanceInterface if you do not need any custom user errors.
/// To work with custom user errors or interpreter internals, use CustomModuleInstanceInterface.
pub type DefaultModuleInstanceInterface = self::module::ModuleInstanceInterface<DummyUserError>;
<commit_msg>from user error for error<commit_after>//! WebAssembly interpreter module.

/// Custom user error.
pub trait UserError: 'static + ::std::fmt::Display + ::std::fmt::Debug + Clone + PartialEq {
}

/// Internal interpreter error.
#[derive(Debug, Clone, PartialEq)]
pub enum Error<E> where E: UserError {
	/// Program-level error.
	Program(String),
	/// Validation error.
	Validation(String),
	/// Initialization error.
	Initialization(String),
	/// Function-level error.
	Function(String),
	/// Table-level error.
	Table(String),
	/// Memory-level error.
	Memory(String),
	/// Variable-level error.
	Variable(String),
	/// Global-level error.
	Global(String),
	/// Local-level error.
	Local(String),
	/// Stack-level error.
	Stack(String),
	/// Value-level error.
	Value(String),
	/// Interpreter (code) error.
	Interpreter(String),
	/// Env module error.
	Env(String),
	/// Native module error.
	Native(String),
	/// Trap.
	Trap(String),
	/// Custom user error.
	User(E),
}

impl<E> Into<String> for Error<E> where E: UserError {
	fn into(self) -> String {
		match self {
			Error::Program(s) => s,
			Error::Validation(s) => s,
			Error::Initialization(s) => s,
			Error::Function(s) => s,
			Error::Table(s) => s,
			Error::Memory(s) => s,
			Error::Variable(s) => s,
			Error::Global(s) => s,
			Error::Local(s) => s,
			Error::Stack(s) => s,
			Error::Interpreter(s) => s,
			Error::Value(s) => s,
			Error::Env(s) => s,
			Error::Native(s) => s,
			Error::Trap(s) => format!("trap: {}", s),
			Error::User(e) => format!("user: {}", e),
		}
	}
}

/// Dummy user error.
#[derive(Debug, Clone, PartialEq)]
pub struct DummyUserError;

impl UserError for DummyUserError {}

impl ::std::fmt::Display for DummyUserError {
	fn fmt(&self, _f: &mut ::std::fmt::Formatter) -> Result<(), ::std::fmt::Error> { Ok(()) }
}

impl<U> From<U> for Error<U> where U: UserError + Sized {
	fn from(e: U) -> Self {
		Error::User(e)
	}
}

mod env;
mod env_native;
mod imports;
mod memory;
mod module;
mod program;
mod runner;
mod stack;
mod table;
mod validator;
mod value;
mod variable;

#[cfg(test)]
mod tests;

pub use self::memory::MemoryInstance;
pub use self::module::{ModuleInstance, ModuleInstanceInterface,
	ItemIndex, ExportEntryType, CallerContext, ExecutionParams, FunctionSignature};
pub use self::table::TableInstance;
pub use self::program::ProgramInstance;
pub use self::value::RuntimeValue;
pub use self::variable::{VariableInstance, VariableType, ExternalVariableValue};
pub use self::env_native::{env_native_module, UserDefinedElements, UserFunctionExecutor, UserFunctionDescriptor};
pub use self::env::EnvParams;

/// Default type of Error if you do not need any custom user errors.
pub type DummyError = Error<DummyUserError>;

/// Default type of ProgramInstance if you do not need any custom user errors.
/// To work with custom user errors or interpreter internals, use CustomProgramInstance.
pub type DefaultProgramInstance = self::program::ProgramInstance<DummyUserError>;

/// Default type of ModuleInstance if you do not need any custom user errors.
/// To work with custom user errors or interpreter internals, use CustomModuleInstance.
pub type DefaultModuleInstance = self::module::ModuleInstance<DummyUserError>;

/// Default type of ModuleInstanceInterface if you do not need any custom user errors.
/// To work with custom user errors or interpreter internals, use CustomModuleInstanceInterface.
pub type DefaultModuleInstanceInterface = self::module::ModuleInstanceInterface<DummyUserError>;
<|endoftext|>


<commit_before>// TODO rm allow(unused)
#![allow(unsafe_code)]

use std::{
    convert::{TryFrom, TryInto},
    fmt::{self, Debug},
    fs::File,
    mem::{transmute, MaybeUninit},
    path::Path,
    sync::{
        atomic::{AtomicU32, Ordering::Acquire},
        Arc,
    },
};

use crossbeam_epoch::pin;

use crate::{
    pagecache::{pread_exact, pwrite_all, MessageKind},
    stack::Stack,
    Error, Lsn, Result,
};

#[cfg(not(feature = "testing"))]
const MIN_SZ: u64 = 64 * 1024;

#[cfg(feature = "testing")]
const MIN_SZ: u64 = 32;

const MIN_TRAILING_ZEROS: u64 = MIN_SZ.trailing_zeros() as u64;

pub type SlabId = u8;
pub type SlabIdx = u32;

/// A unique identifier for a particular slot in the heap
#[allow(clippy::module_name_repetitions)]
#[derive(Clone, Copy, PartialOrd, Ord, Eq, PartialEq, Hash)]
pub struct HeapId {
    pub location: u64,
    pub original_lsn: Lsn,
}

impl Debug for HeapId {
    fn fmt(
        &self,
        f: &mut fmt::Formatter<'_>,
    ) -> std::result::Result<(), fmt::Error> {
        let (slab, idx, original_lsn) = self.decompose();
        f.debug_struct("HeapId")
            .field("slab", &slab)
            .field("idx", &idx)
            .field("original_lsn", &original_lsn)
            .finish()
    }
}

impl HeapId {
    pub fn decompose(&self) -> (SlabId, SlabIdx, Lsn) {
        const IDX_MASK: u64 = (1 << 32) - 1;
        let slab_id =
            u8::try_from((self.location >> 32).trailing_zeros()).unwrap();
        let slab_idx = u32::try_from(self.location & IDX_MASK).unwrap();
        (slab_id, slab_idx, self.original_lsn)
    }

    pub fn compose(
        slab_id: SlabId,
        slab_idx: SlabIdx,
        original_lsn: Lsn,
    ) -> HeapId {
        let slab = 1 << (32 + u64::from(slab_id));
        let heap_id = slab | u64::from(slab_idx);
        HeapId { location: heap_id, original_lsn }
    }

    fn offset(&self) -> u64 {
        let (slab_id, idx, _) = self.decompose();
        slab_id_to_size(slab_id) * u64::from(idx)
    }

    fn slab_size(&self) -> u64 {
        let (slab_id, _idx, _lsn) = self.decompose();
        slab_id_to_size(slab_id)
    }
}

pub(crate) fn slab_size(size: u64) -> u64 {
    slab_id_to_size(size_to_slab_id(size))
}

fn slab_id_to_size(slab_id: u8) -> u64 {
    1 << (MIN_TRAILING_ZEROS + u64::from(slab_id))
}

fn size_to_slab_id(size: u64) -> SlabId {
    // find the power of 2 that is at least 64k
    let normalized_size = std::cmp::max(MIN_SZ, size.next_power_of_two());

    // drop the lowest unused bits
    let rebased_size = normalized_size >> MIN_TRAILING_ZEROS;

    u8::try_from(rebased_size.trailing_zeros()).unwrap()
}

pub(crate) struct Reservation {
    slab_free: Arc<Stack<u32>>,
    completed: bool,
    file: File,
    pub heap_id: HeapId,
    from_tip: bool,
}

impl Drop for Reservation {
    fn drop(&mut self) {
        if !self.completed {
            let (_slab_id, idx, _) = self.heap_id.decompose();
            self.slab_free.push(idx, &pin());
        }
    }
}

impl Reservation {
    pub fn complete(mut self, data: &[u8]) -> Result<HeapId> {
        log::trace!(
            "Heap::complete({:?}) to offset {} in file {:?}",
            self.heap_id,
            self.heap_id.offset(),
            self.file
        );
        assert_eq!(data.len() as u64, self.heap_id.slab_size());

        // write data
        pwrite_all(&self.file, data, self.heap_id.offset())?;

        // sync data
        if self.from_tip {
            self.file.sync_all()?;
        } else if cfg!(not(target_os = "linux")) {
            self.file.sync_data()?;
        } else {
            #[allow(clippy::assertions_on_constants)]
            {
                assert!(cfg!(target_os = "linux"));
            }

            #[cfg(target_os = "linux")]
            {
                use std::os::unix::io::AsRawFd;
                let ret = unsafe {
                    libc::sync_file_range(
                        self.file.as_raw_fd(),
                        i64::try_from(self.heap_id.offset()).unwrap(),
                        i64::try_from(data.len()).unwrap(),
                        libc::SYNC_FILE_RANGE_WAIT_BEFORE
                            | libc::SYNC_FILE_RANGE_WRITE
                            | libc::SYNC_FILE_RANGE_WAIT_AFTER,
                    )
                };
                if ret < 0 {
                    let err = std::io::Error::last_os_error();
                    if let Some(libc::ENOSYS) = err.raw_os_error() {
                        self.file.sync_all()?;
                    } else {
                        return Err(err.into());
                    }
                }
            }
        }

        // if this is not reached due to an IO error,
        // the offset will be returned to the Slab in Drop
        self.completed = true;

        Ok(self.heap_id)
    }
}

#[derive(Debug)]
pub(crate) struct Heap {
    // each slab stores
    // items that are double
    // the size of the previous,
    // ranging from 64k in the
    // smallest slab to 2^48 in
    // the last.
    slabs: [Slab; 32],
}

impl Heap {
    pub fn start<P: AsRef<Path>>(p: P) -> Result<Heap> {
        let mut slabs: [MaybeUninit<Slab>; 32] = unsafe { std::mem::zeroed() };

        for slab_id in 0..32 {
            let slab = Slab::start(&p, slab_id)?;
            slabs[slab_id as usize] = MaybeUninit::new(slab);
        }

        Ok(Heap { slabs: unsafe { transmute(slabs) } })
    }

    pub fn gc_unknown_items(&self, snapshot: &crate::pagecache::Snapshot) {
        let mut bitmaps = vec![];
        for slab in &self.slabs {
            let tip = slab.tip.load(Acquire) as usize;
            bitmaps.push(vec![0_u64; 1 + (tip / 64)]);
        }

        for page_state in &snapshot.pt {
            for heap_id in page_state.heap_ids() {
                let (slab_id, idx, _lsn) = heap_id.decompose();

                // set the bit for this slot
                let block = idx / 64;
                let bit = idx % 64;
                let bitmask = 1 << bit;
                bitmaps[slab_id as usize][block as usize] |= bitmask;
            }
        }

        let iter = self.slabs.iter().zip(bitmaps.into_iter());

        for (slab, bitmap) in iter {
            let tip = slab.tip.load(Acquire);

            for idx in 0..tip {
                let block = idx / 64;
                let bit = idx % 64;
                let bitmask = 1 << bit;
                let free = bitmap[block as usize] & bitmask == 0;

                if free {
                    slab.free(idx);
                }
            }
        }
    }

    pub fn read(
        &self,
        heap_id: HeapId,
        use_compression: bool,
    ) -> Result<(MessageKind, Vec<u8>)> {
        log::trace!("Heap::read({:?})", heap_id);
        let (slab_id, slab_idx, original_lsn) = heap_id.decompose();
        self.slabs[slab_id as usize].read(
            slab_idx,
            original_lsn,
            use_compression,
        )
    }

    pub fn free(&self, heap_id: HeapId) {
        log::trace!("Heap::free({:?})", heap_id);
        let (slab_id, slab_idx, _) = heap_id.decompose();
        self.slabs[slab_id as usize].free(slab_idx)
    }

    pub fn reserve(&self, size: u64, original_lsn: Lsn) -> Reservation {
        assert!(size < 1 << 48);
        let slab_id = size_to_slab_id(size);
        let ret = self.slabs[slab_id as usize].reserve(original_lsn);
        log::trace!("Heap::reserve({}) -> {:?}", size, ret.heap_id);
        ret
    }
}

#[derive(Debug)]
struct Slab {
    file: File,
    slab_id: u8,
    tip: AtomicU32,
    free: Arc<Stack<u32>>,
}

impl Slab {
    pub fn start<P: AsRef<Path>>(directory: P, slab_id: u8) -> Result<Slab> {
        let bs = slab_id_to_size(slab_id);
        let free = Arc::new(Stack::default());

        let mut options = std::fs::OpenOptions::new();
        options.create(true);
        options.read(true);
        options.write(true);

        let file =
            options.open(directory.as_ref().join(format!("{:02}", slab_id)))?;
        let len = file.metadata()?.len();
        let max_idx = len / bs;
        log::trace!(
            "starting heap slab for sizes of {}. tip: {} max idx: {}",
            bs,
            len,
            max_idx
        );
        let tip = AtomicU32::new(u32::try_from(max_idx).unwrap());

        Ok(Slab { file, slab_id, tip, free })
    }

    fn read(
        &self,
        slab_idx: SlabIdx,
        original_lsn: Lsn,
        use_compression: bool,
    ) -> Result<(MessageKind, Vec<u8>)> {
        let bs = slab_id_to_size(self.slab_id);
        let offset = u64::from(slab_idx) * bs;

        log::trace!("reading heap slab slot {} at offset {}", slab_idx, offset);

        let mut heap_buf = vec![0; usize::try_from(bs).unwrap()];

        pread_exact(&self.file, &mut heap_buf, offset)?;

        let stored_crc =
            u32::from_le_bytes(heap_buf[1..5].as_ref().try_into().unwrap());

        let mut hasher = crc32fast::Hasher::new();
        hasher.update(&heap_buf[0..1]);
        hasher.update(&heap_buf[5..]);
        let actual_crc = hasher.finalize();

        if actual_crc == stored_crc {
            let actual_lsn = Lsn::from_le_bytes(
                heap_buf[5..13].as_ref().try_into().unwrap(),
            );
            if actual_lsn != original_lsn {
                log::error!(
                    "heap slot lsn {} does not match expected original lsn {}",
                    actual_lsn,
                    original_lsn
                );
                return Err(Error::corruption(None));
            }
            let buf = heap_buf[13..].to_vec();
            let buf = if use_compression {
                crate::pagecache::decompress(buf)
            } else {
                buf
            };
            Ok((MessageKind::from(heap_buf[0]), buf))
        } else {
            log::error!(
                "heap message CRC does not match contents. stored: {} actual: {}",
                stored_crc,
                actual_crc
            );
            Err(Error::corruption(None))
        }
    }

    fn reserve(&self, original_lsn: Lsn) -> Reservation {
        let (idx, from_tip) = if let Some(idx) = self.free.pop(&pin()) {
            log::trace!(
                "reusing heap index {} in slab for sizes of {}",
                idx,
                slab_id_to_size(self.slab_id),
            );
            (idx, false)
        } else {
            log::trace!(
                "no free heap slots in slab for sizes of {}",
                slab_id_to_size(self.slab_id),
            );
            (self.tip.fetch_add(1, Acquire), true)
        };

        log::trace!(
            "heap reservation for slot {} in the slab for sizes of {}",
            idx,
            slab_id_to_size(self.slab_id),
        );

        let heap_id = HeapId::compose(self.slab_id, idx, original_lsn);

        Reservation {
            slab_free: self.free.clone(),
            completed: false,
            file: self.file.try_clone().unwrap(),
            from_tip,
            heap_id,
        }
    }

    fn free(&self, idx: u32) {
        self.punch_hole(idx);
        self.free.push(idx, &pin());
    }

    fn punch_hole(&self, #[allow(unused)] idx: u32) {
        #[cfg(all(target_os = "linux", not(miri)))]
        {
            use std::{
                os::unix::io::AsRawFd,
                sync::atomic::{AtomicBool, Ordering::Relaxed},
            };

            use libc::{fallocate, FALLOC_FL_KEEP_SIZE, FALLOC_FL_PUNCH_HOLE};

            static HOLE_PUNCHING_ENABLED: AtomicBool = AtomicBool::new(true);
            const MODE: i32 = FALLOC_FL_KEEP_SIZE | FALLOC_FL_PUNCH_HOLE;

            if HOLE_PUNCHING_ENABLED.load(Relaxed) {
                let bs = i64::try_from(slab_id_to_size(self.slab_id)).unwrap();
                let offset = i64::from(idx) * bs;

                let fd = self.file.as_raw_fd();

                let ret = unsafe {
                    fallocate(
                        fd,
                        MODE,
                        #[allow(clippy::useless_conversion)]
                        offset.try_into().unwrap(),
                        #[allow(clippy::useless_conversion)]
                        bs.try_into().unwrap(),
                    )
                };

                if ret != 0 {
                    let err = std::io::Error::last_os_error();
                    log::error!(
                        "failed to punch hole in heap file: {:?}. disabling hole punching",
                        err
                    );
                    HOLE_PUNCHING_ENABLED.store(false, Relaxed);
                }
            }
        }
    }
}
<commit_msg>Reduce log severity for heap writes<commit_after>// TODO rm allow(unused)
#![allow(unsafe_code)]

use std::{
    convert::{TryFrom, TryInto},
    fmt::{self, Debug},
    fs::File,
    mem::{transmute, MaybeUninit},
    path::Path,
    sync::{
        atomic::{AtomicU32, Ordering::Acquire},
        Arc,
    },
};

use crossbeam_epoch::pin;

use crate::{
    pagecache::{pread_exact, pwrite_all, MessageKind},
    stack::Stack,
    Error, Lsn, Result,
};

#[cfg(not(feature = "testing"))]
const MIN_SZ: u64 = 64 * 1024;

#[cfg(feature = "testing")]
const MIN_SZ: u64 = 32;

const MIN_TRAILING_ZEROS: u64 = MIN_SZ.trailing_zeros() as u64;

pub type SlabId = u8;
pub type SlabIdx = u32;

/// A unique identifier for a particular slot in the heap
#[allow(clippy::module_name_repetitions)]
#[derive(Clone, Copy, PartialOrd, Ord, Eq, PartialEq, Hash)]
pub struct HeapId {
    pub location: u64,
    pub original_lsn: Lsn,
}

impl Debug for HeapId {
    fn fmt(
        &self,
        f: &mut fmt::Formatter<'_>,
    ) -> std::result::Result<(), fmt::Error> {
        let (slab, idx, original_lsn) = self.decompose();
        f.debug_struct("HeapId")
            .field("slab", &slab)
            .field("idx", &idx)
            .field("original_lsn", &original_lsn)
            .finish()
    }
}

impl HeapId {
    pub fn decompose(&self) -> (SlabId, SlabIdx, Lsn) {
        const IDX_MASK: u64 = (1 << 32) - 1;
        let slab_id =
            u8::try_from((self.location >> 32).trailing_zeros()).unwrap();
        let slab_idx = u32::try_from(self.location & IDX_MASK).unwrap();
        (slab_id, slab_idx, self.original_lsn)
    }

    pub fn compose(
        slab_id: SlabId,
        slab_idx: SlabIdx,
        original_lsn: Lsn,
    ) -> HeapId {
        let slab = 1 << (32 + u64::from(slab_id));
        let heap_id = slab | u64::from(slab_idx);
        HeapId { location: heap_id, original_lsn }
    }

    fn offset(&self) -> u64 {
        let (slab_id, idx, _) = self.decompose();
        slab_id_to_size(slab_id) * u64::from(idx)
    }

    fn slab_size(&self) -> u64 {
        let (slab_id, _idx, _lsn) = self.decompose();
        slab_id_to_size(slab_id)
    }
}

pub(crate) fn slab_size(size: u64) -> u64 {
    slab_id_to_size(size_to_slab_id(size))
}

fn slab_id_to_size(slab_id: u8) -> u64 {
    1 << (MIN_TRAILING_ZEROS + u64::from(slab_id))
}

fn size_to_slab_id(size: u64) -> SlabId {
    // find the power of 2 that is at least 64k
    let normalized_size = std::cmp::max(MIN_SZ, size.next_power_of_two());

    // drop the lowest unused bits
    let rebased_size = normalized_size >> MIN_TRAILING_ZEROS;

    u8::try_from(rebased_size.trailing_zeros()).unwrap()
}

pub(crate) struct Reservation {
    slab_free: Arc<Stack<u32>>,
    completed: bool,
    file: File,
    pub heap_id: HeapId,
    from_tip: bool,
}

impl Drop for Reservation {
    fn drop(&mut self) {
        if !self.completed {
            let (_slab_id, idx, _) = self.heap_id.decompose();
            self.slab_free.push(idx, &pin());
        }
    }
}

impl Reservation {
    pub fn complete(mut self, data: &[u8]) -> Result<HeapId> {
        log::trace!(
            "Heap::complete({:?}) to offset {} in file {:?}",
            self.heap_id,
            self.heap_id.offset(),
            self.file
        );
        assert_eq!(data.len() as u64, self.heap_id.slab_size());

        // write data
        pwrite_all(&self.file, data, self.heap_id.offset())?;

        // sync data
        if self.from_tip {
            self.file.sync_all()?;
        } else if cfg!(not(target_os = "linux")) {
            self.file.sync_data()?;
        } else {
            #[allow(clippy::assertions_on_constants)]
            {
                assert!(cfg!(target_os = "linux"));
            }

            #[cfg(target_os = "linux")]
            {
                use std::os::unix::io::AsRawFd;
                let ret = unsafe {
                    libc::sync_file_range(
                        self.file.as_raw_fd(),
                        i64::try_from(self.heap_id.offset()).unwrap(),
                        i64::try_from(data.len()).unwrap(),
                        libc::SYNC_FILE_RANGE_WAIT_BEFORE
                            | libc::SYNC_FILE_RANGE_WRITE
                            | libc::SYNC_FILE_RANGE_WAIT_AFTER,
                    )
                };
                if ret < 0 {
                    let err = std::io::Error::last_os_error();
                    if let Some(libc::ENOSYS) = err.raw_os_error() {
                        self.file.sync_all()?;
                    } else {
                        return Err(err.into());
                    }
                }
            }
        }

        // if this is not reached due to an IO error,
        // the offset will be returned to the Slab in Drop
        self.completed = true;

        Ok(self.heap_id)
    }
}

#[derive(Debug)]
pub(crate) struct Heap {
    // each slab stores
    // items that are double
    // the size of the previous,
    // ranging from 64k in the
    // smallest slab to 2^48 in
    // the last.
    slabs: [Slab; 32],
}

impl Heap {
    pub fn start<P: AsRef<Path>>(p: P) -> Result<Heap> {
        let mut slabs: [MaybeUninit<Slab>; 32] = unsafe { std::mem::zeroed() };

        for slab_id in 0..32 {
            let slab = Slab::start(&p, slab_id)?;
            slabs[slab_id as usize] = MaybeUninit::new(slab);
        }

        Ok(Heap { slabs: unsafe { transmute(slabs) } })
    }

    pub fn gc_unknown_items(&self, snapshot: &crate::pagecache::Snapshot) {
        let mut bitmaps = vec![];
        for slab in &self.slabs {
            let tip = slab.tip.load(Acquire) as usize;
            bitmaps.push(vec![0_u64; 1 + (tip / 64)]);
        }

        for page_state in &snapshot.pt {
            for heap_id in page_state.heap_ids() {
                let (slab_id, idx, _lsn) = heap_id.decompose();

                // set the bit for this slot
                let block = idx / 64;
                let bit = idx % 64;
                let bitmask = 1 << bit;
                bitmaps[slab_id as usize][block as usize] |= bitmask;
            }
        }

        let iter = self.slabs.iter().zip(bitmaps.into_iter());

        for (slab, bitmap) in iter {
            let tip = slab.tip.load(Acquire);

            for idx in 0..tip {
                let block = idx / 64;
                let bit = idx % 64;
                let bitmask = 1 << bit;
                let free = bitmap[block as usize] & bitmask == 0;

                if free {
                    slab.free(idx);
                }
            }
        }
    }

    pub fn read(
        &self,
        heap_id: HeapId,
        use_compression: bool,
    ) -> Result<(MessageKind, Vec<u8>)> {
        log::trace!("Heap::read({:?})", heap_id);
        let (slab_id, slab_idx, original_lsn) = heap_id.decompose();
        self.slabs[slab_id as usize].read(
            slab_idx,
            original_lsn,
            use_compression,
        )
    }

    pub fn free(&self, heap_id: HeapId) {
        log::trace!("Heap::free({:?})", heap_id);
        let (slab_id, slab_idx, _) = heap_id.decompose();
        self.slabs[slab_id as usize].free(slab_idx)
    }

    pub fn reserve(&self, size: u64, original_lsn: Lsn) -> Reservation {
        assert!(size < 1 << 48);
        let slab_id = size_to_slab_id(size);
        let ret = self.slabs[slab_id as usize].reserve(original_lsn);
        log::trace!("Heap::reserve({}) -> {:?}", size, ret.heap_id);
        ret
    }
}

#[derive(Debug)]
struct Slab {
    file: File,
    slab_id: u8,
    tip: AtomicU32,
    free: Arc<Stack<u32>>,
}

impl Slab {
    pub fn start<P: AsRef<Path>>(directory: P, slab_id: u8) -> Result<Slab> {
        let bs = slab_id_to_size(slab_id);
        let free = Arc::new(Stack::default());

        let mut options = std::fs::OpenOptions::new();
        options.create(true);
        options.read(true);
        options.write(true);

        let file =
            options.open(directory.as_ref().join(format!("{:02}", slab_id)))?;
        let len = file.metadata()?.len();
        let max_idx = len / bs;
        log::trace!(
            "starting heap slab for sizes of {}. tip: {} max idx: {}",
            bs,
            len,
            max_idx
        );
        let tip = AtomicU32::new(u32::try_from(max_idx).unwrap());

        Ok(Slab { file, slab_id, tip, free })
    }

    fn read(
        &self,
        slab_idx: SlabIdx,
        original_lsn: Lsn,
        use_compression: bool,
    ) -> Result<(MessageKind, Vec<u8>)> {
        let bs = slab_id_to_size(self.slab_id);
        let offset = u64::from(slab_idx) * bs;

        log::trace!("reading heap slab slot {} at offset {}", slab_idx, offset);

        let mut heap_buf = vec![0; usize::try_from(bs).unwrap()];

        pread_exact(&self.file, &mut heap_buf, offset)?;

        let stored_crc =
            u32::from_le_bytes(heap_buf[1..5].as_ref().try_into().unwrap());

        let mut hasher = crc32fast::Hasher::new();
        hasher.update(&heap_buf[0..1]);
        hasher.update(&heap_buf[5..]);
        let actual_crc = hasher.finalize();

        if actual_crc == stored_crc {
            let actual_lsn = Lsn::from_le_bytes(
                heap_buf[5..13].as_ref().try_into().unwrap(),
            );
            if actual_lsn != original_lsn {
                log::debug!(
                    "heap slot lsn {} does not match expected original lsn {}",
                    actual_lsn,
                    original_lsn
                );
                return Err(Error::corruption(None));
            }
            let buf = heap_buf[13..].to_vec();
            let buf = if use_compression {
                crate::pagecache::decompress(buf)
            } else {
                buf
            };
            Ok((MessageKind::from(heap_buf[0]), buf))
        } else {
            log::debug!(
                "heap message CRC does not match contents. stored: {} actual: {}",
                stored_crc,
                actual_crc
            );
            Err(Error::corruption(None))
        }
    }

    fn reserve(&self, original_lsn: Lsn) -> Reservation {
        let (idx, from_tip) = if let Some(idx) = self.free.pop(&pin()) {
            log::trace!(
                "reusing heap index {} in slab for sizes of {}",
                idx,
                slab_id_to_size(self.slab_id),
            );
            (idx, false)
        } else {
            log::trace!(
                "no free heap slots in slab for sizes of {}",
                slab_id_to_size(self.slab_id),
            );
            (self.tip.fetch_add(1, Acquire), true)
        };

        log::trace!(
            "heap reservation for slot {} in the slab for sizes of {}",
            idx,
            slab_id_to_size(self.slab_id),
        );

        let heap_id = HeapId::compose(self.slab_id, idx, original_lsn);

        Reservation {
            slab_free: self.free.clone(),
            completed: false,
            file: self.file.try_clone().unwrap(),
            from_tip,
            heap_id,
        }
    }

    fn free(&self, idx: u32) {
        self.punch_hole(idx);
        self.free.push(idx, &pin());
    }

    fn punch_hole(&self, #[allow(unused)] idx: u32) {
        #[cfg(all(target_os = "linux", not(miri)))]
        {
            use std::{
                os::unix::io::AsRawFd,
                sync::atomic::{AtomicBool, Ordering::Relaxed},
            };

            use libc::{fallocate, FALLOC_FL_KEEP_SIZE, FALLOC_FL_PUNCH_HOLE};

            static HOLE_PUNCHING_ENABLED: AtomicBool = AtomicBool::new(true);
            const MODE: i32 = FALLOC_FL_KEEP_SIZE | FALLOC_FL_PUNCH_HOLE;

            if HOLE_PUNCHING_ENABLED.load(Relaxed) {
                let bs = i64::try_from(slab_id_to_size(self.slab_id)).unwrap();
                let offset = i64::from(idx) * bs;

                let fd = self.file.as_raw_fd();

                let ret = unsafe {
                    fallocate(
                        fd,
                        MODE,
                        #[allow(clippy::useless_conversion)]
                        offset.try_into().unwrap(),
                        #[allow(clippy::useless_conversion)]
                        bs.try_into().unwrap(),
                    )
                };

                if ret != 0 {
                    let err = std::io::Error::last_os_error();
                    log::error!(
                        "failed to punch hole in heap file: {:?}. disabling hole punching",
                        err
                    );
                    HOLE_PUNCHING_ENABLED.store(false, Relaxed);
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::rc::Rc;
use super::constants::*;
use super::nodes::*;

/// Store a command tree while populating it. This is used
/// to construct a [`RootNode`] to be used with the [`Parser`].
///
/// The lifetime parameter `'a` refers to the lifetime
/// of the strings used for [command] and [parameter] names and
/// help text.
///
/// [command]: struct.Command.html
/// [parameter]: struct.Parameter.html
/// [`Parser`]: struct.Parser.html
/// [`RootNode`]: struct.RootNode.html
pub struct CommandTree<'a> {
    commands: Vec<Command<'a>>,
}

impl<'a> Default for CommandTree<'a> {
    fn default() -> Self {
        CommandTree { commands: vec![] }
    }
}

impl<'a> CommandTree<'a> {
    /// Create a new `CommandTree`.
    pub fn new() -> Self {
        Default::default()
    }

    /// Add a `Command` to the `CommandTree`.
    pub fn command(&mut self, command: Command<'a>) {
        self.commands.push(command);
    }

    /// Construct the `CommandTree` and produce a `RootNode`.
    pub fn finalize(&self) -> Rc<Node> {
        let mut successors: Vec<Rc<Node>> = vec![];
        for c in &self.commands {
            successors.push(Rc::new(Node::Command(self.build_command(c))));
        }
        Rc::new(Node::Root(RootNode::new(successors)))
    }

    fn build_command(&self, command: &Command) -> CommandNode {
        let mut parameters: Vec<Rc<Node>> = vec![];
        let mut successors: Vec<Rc<Node>> = vec![];
        for parameter in &command.parameters {
            match parameter.kind {
                ParameterKind::Flag => {
                    self.build_flag_parameter(parameter, &mut parameters, &mut successors);
                }
                ParameterKind::Named => {
                    self.build_named_parameter(parameter, &mut parameters, &mut successors);
                }
                ParameterKind::Simple => {
                    self.build_simple_parameter(parameter, &mut parameters, &mut successors);
                }
            };
        }
        // We'll want to find the right node for the wrapped_root
        // and pass it along here.
        CommandNode::new(command.name,
                         command.help_text,
                         command.hidden,
                         command.priority,
                         successors,
                         None,
                         parameters)
    }

    fn build_flag_parameter(&self,
                            parameter: &Parameter,
                            parameters: &mut Vec<Rc<Node>>,
                            successors: &mut Vec<Rc<Node>>) {
        let p = ParameterNode::new(parameter.name,
                                   parameter.help_text,
                                   parameter.hidden,
                                   parameter.priority.unwrap_or(PRIORITY_DEFAULT),
                                   vec![],
                                   parameter.repeatable,
                                   None,
                                   parameter.kind,
                                   parameter.required);
        let p = Rc::new(Node::Parameter(p));
        parameters.push(p.clone());
        successors.push(p);
    }

    fn build_named_parameter(&self,
                             parameter: &Parameter,
                             parameters: &mut Vec<Rc<Node>>,
                             successors: &mut Vec<Rc<Node>>) {
        let p = ParameterNode::new(parameter.name,
                                   parameter.help_text,
                                   parameter.hidden,
                                   parameter.priority
                                            .unwrap_or(PRIORITY_PARAMETER),
                                   vec![],
                                   parameter.repeatable,
                                   None,
                                   parameter.kind,
                                   parameter.required);
        let p = Rc::new(Node::Parameter(p));
        parameters.push(p.clone());
        let n = ParameterNameNode::new(parameter.name,
                                       parameter.hidden,
                                       PRIORITY_DEFAULT,
                                       vec![p.clone()],
                                       parameter.repeatable,
                                       Some(p.clone()),
                                       p.clone());
        successors.push(Rc::new(Node::ParameterName(n)));
        for alias in &parameter.aliases {
            let a = ParameterNameNode::new(alias,
                                           parameter.hidden,
                                           PRIORITY_DEFAULT,
                                           vec![p.clone()],
                                           parameter.repeatable,
                                           Some(p.clone()),
                                           p.clone());
            successors.push(Rc::new(Node::ParameterName(a)));
        }
    }

    fn build_simple_parameter(&self,
                              parameter: &Parameter,
                              parameters: &mut Vec<Rc<Node>>,
                              successors: &mut Vec<Rc<Node>>) {
        let p = ParameterNode::new(parameter.name,
                                   parameter.help_text,
                                   parameter.hidden,
                                   parameter.priority.unwrap_or(PRIORITY_PARAMETER),
                                   vec![],
                                   parameter.repeatable,
                                   None,
                                   parameter.kind,
                                   parameter.required);
        let p = Rc::new(Node::Parameter(p));
        parameters.push(p.clone());
        successors.push(p.clone());
    }
}

/// Description of a command to be added to the [`CommandTree`].
///
/// The lifetime parameter `'a` refers to the lifetime
/// of the strings used for command names and help text.
///
/// [`CommandTree`]: struct.CommandTree.html
pub struct Command<'a> {
    hidden: bool,
    priority: i32,
    name: &'a str,
    help_text: Option<&'a str>,
    parameters: Vec<Parameter<'a>>,
    wrapped_root: Option<String>,
}

impl<'a> Command<'a> {
    /// Construct a default (blank) command with the given `name`.
    pub fn new(name: &'a str) -> Self {
        Command {
            hidden: false,
            priority: PRIORITY_DEFAULT,
            name: name,
            help_text: None,
            parameters: vec![],
            wrapped_root: None,
        }
    }

    /// Mark the command as hidden. Hidden commands will match
    /// within the parser, but are not listed during completion.
    pub fn hidden(mut self, hidden: bool) -> Self {
        self.hidden = hidden;
        self
    }

    /// Give the command a priority. This is used when sorting
    /// out conflicts during matching and completion.
    pub fn priority(mut self, priority: i32) -> Self {
        self.priority = priority;
        self
    }

    /// Supply help text for the command.
    pub fn help(mut self, help_text: &'a str) -> Self {
        self.help_text = Some(help_text);
        self
    }

    /// Add a [`Parameter`] to the command.
    ///
    /// [`Parameter`]: struct.Parameter.html
    pub fn parameter(mut self, parameter: Parameter<'a>) -> Self {
        self.parameters.push(parameter);
        self
    }

    /// Create a [`WrapperNode`] instead of a [`CommandNode`]. The
    /// `wrapped_root` signifies the path to the command that should
    /// be wrapped by this command.
    ///
    /// [`CommandNode`]: struct.CommandNode.html
    /// [`WrapperNode`]: struct.WrapperNode.html
    pub fn wraps(mut self, wrapped_root: String) -> Self {
        self.wrapped_root = Some(wrapped_root);
        self
    }
}

/// Description of a parameter to be added to the [`Command`].
///
/// The lifetime parameter `'a` refers to the lifetime
/// of the strings used for parameter names, aliases and
/// help text.
///
/// [`Command`]: struct.Command.html
pub struct Parameter<'a> {
    hidden: bool,
    priority: Option<i32>,
    name: &'a str,
    repeatable: bool,
    aliases: Vec<&'a str>,
    help_text: Option<&'a str>,
    kind: ParameterKind,
    required: bool,
}

impl<'a> Parameter<'a> {
    /// Construct a default (blank) parameter with the given `name`.
    pub fn new(name: &'a str) -> Self {
        Parameter {
            hidden: false,
            priority: None,
            name: name,
            repeatable: false,
            aliases: vec![],
            help_text: None,
            kind: ParameterKind::Simple,
            required: false,
        }
    }

    /// Mark the parameter as hidden. Hidden parameters will match
    /// within the parser, but are not listed during completion.
    pub fn hidden(mut self, hidden: bool) -> Self {
        self.hidden = hidden;
        self
    }

    /// Give the parameter a priority. This is used when sorting
    /// out conflicts during matching and completion.
    ///
    /// The `priority` of a `Parameter` defaults to `PRIORITY_PARAMETER`
    /// except for when the `kind` is `ParameterKind::Flag` in which
    /// case, the default will be `PRIORITY_DEFAULT`.
    pub fn priority(mut self, priority: i32) -> Self {
        self.priority = Some(priority);
        self
    }

    /// Establish whether or not this parameter is repeatable.
    /// Repeated parameters produce a vector of values and can
    /// be given multiple times within a single command invocation.
    pub fn repeatable(mut self, repeatable: bool) -> Self {
        self.repeatable = repeatable;
        self
    }

    /// Add an alias that this parameter can use.
    ///
    /// Aliases are currently only valid for parameters of `kind`
    /// `ParameterKind::Named`.
    pub fn alias(mut self, alias: &'a str) -> Self {
        self.aliases.push(alias);
        self
    }

    /// Supply the help text for the parameter.
    pub fn help(mut self, help_text: &'a str) -> Self {
        self.help_text = Some(help_text);
        self
    }

    /// Establish whether or not this parameter is required.
    pub fn required(mut self, required: bool) -> Self {
        self.required = required;
        self
    }

    /// Set which type of [`ParameterNode`] is supposed to be created
    /// to represent this parameter.
    ///
    /// [`ParameterNode`]: trait.ParameterNode.html
    pub fn kind(mut self, kind: ParameterKind) -> Self {
        self.kind = kind;
        self
    }
}
<commit_msg>Remove reference to the defunct WrapperNode.<commit_after>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::rc::Rc;
use super::constants::*;
use super::nodes::*;

/// Store a command tree while populating it. This is used
/// to construct a [`RootNode`] to be used with the [`Parser`].
///
/// The lifetime parameter `'a` refers to the lifetime
/// of the strings used for [command] and [parameter] names and
/// help text.
///
/// [command]: struct.Command.html
/// [parameter]: struct.Parameter.html
/// [`Parser`]: struct.Parser.html
/// [`RootNode`]: struct.RootNode.html
pub struct CommandTree<'a> {
    commands: Vec<Command<'a>>,
}

impl<'a> Default for CommandTree<'a> {
    fn default() -> Self {
        CommandTree { commands: vec![] }
    }
}

impl<'a> CommandTree<'a> {
    /// Create a new `CommandTree`.
    pub fn new() -> Self {
        Default::default()
    }

    /// Add a `Command` to the `CommandTree`.
    pub fn command(&mut self, command: Command<'a>) {
        self.commands.push(command);
    }

    /// Construct the `CommandTree` and produce a `RootNode`.
    pub fn finalize(&self) -> Rc<Node> {
        let mut successors: Vec<Rc<Node>> = vec![];
        for c in &self.commands {
            successors.push(Rc::new(Node::Command(self.build_command(c))));
        }
        Rc::new(Node::Root(RootNode::new(successors)))
    }

    fn build_command(&self, command: &Command) -> CommandNode {
        let mut parameters: Vec<Rc<Node>> = vec![];
        let mut successors: Vec<Rc<Node>> = vec![];
        for parameter in &command.parameters {
            match parameter.kind {
                ParameterKind::Flag => {
                    self.build_flag_parameter(parameter, &mut parameters, &mut successors);
                }
                ParameterKind::Named => {
                    self.build_named_parameter(parameter, &mut parameters, &mut successors);
                }
                ParameterKind::Simple => {
                    self.build_simple_parameter(parameter, &mut parameters, &mut successors);
                }
            };
        }
        // We'll want to find the right node for the wrapped_root
        // and pass it along here.
        CommandNode::new(command.name,
                         command.help_text,
                         command.hidden,
                         command.priority,
                         successors,
                         None,
                         parameters)
    }

    fn build_flag_parameter(&self,
                            parameter: &Parameter,
                            parameters: &mut Vec<Rc<Node>>,
                            successors: &mut Vec<Rc<Node>>) {
        let p = ParameterNode::new(parameter.name,
                                   parameter.help_text,
                                   parameter.hidden,
                                   parameter.priority.unwrap_or(PRIORITY_DEFAULT),
                                   vec![],
                                   parameter.repeatable,
                                   None,
                                   parameter.kind,
                                   parameter.required);
        let p = Rc::new(Node::Parameter(p));
        parameters.push(p.clone());
        successors.push(p);
    }

    fn build_named_parameter(&self,
                             parameter: &Parameter,
                             parameters: &mut Vec<Rc<Node>>,
                             successors: &mut Vec<Rc<Node>>) {
        let p = ParameterNode::new(parameter.name,
                                   parameter.help_text,
                                   parameter.hidden,
                                   parameter.priority
                                            .unwrap_or(PRIORITY_PARAMETER),
                                   vec![],
                                   parameter.repeatable,
                                   None,
                                   parameter.kind,
                                   parameter.required);
        let p = Rc::new(Node::Parameter(p));
        parameters.push(p.clone());
        let n = ParameterNameNode::new(parameter.name,
                                       parameter.hidden,
                                       PRIORITY_DEFAULT,
                                       vec![p.clone()],
                                       parameter.repeatable,
                                       Some(p.clone()),
                                       p.clone());
        successors.push(Rc::new(Node::ParameterName(n)));
        for alias in &parameter.aliases {
            let a = ParameterNameNode::new(alias,
                                           parameter.hidden,
                                           PRIORITY_DEFAULT,
                                           vec![p.clone()],
                                           parameter.repeatable,
                                           Some(p.clone()),
                                           p.clone());
            successors.push(Rc::new(Node::ParameterName(a)));
        }
    }

    fn build_simple_parameter(&self,
                              parameter: &Parameter,
                              parameters: &mut Vec<Rc<Node>>,
                              successors: &mut Vec<Rc<Node>>) {
        let p = ParameterNode::new(parameter.name,
                                   parameter.help_text,
                                   parameter.hidden,
                                   parameter.priority.unwrap_or(PRIORITY_PARAMETER),
                                   vec![],
                                   parameter.repeatable,
                                   None,
                                   parameter.kind,
                                   parameter.required);
        let p = Rc::new(Node::Parameter(p));
        parameters.push(p.clone());
        successors.push(p.clone());
    }
}

/// Description of a command to be added to the [`CommandTree`].
///
/// The lifetime parameter `'a` refers to the lifetime
/// of the strings used for command names and help text.
///
/// [`CommandTree`]: struct.CommandTree.html
pub struct Command<'a> {
    hidden: bool,
    priority: i32,
    name: &'a str,
    help_text: Option<&'a str>,
    parameters: Vec<Parameter<'a>>,
    wrapped_root: Option<String>,
}

impl<'a> Command<'a> {
    /// Construct a default (blank) command with the given `name`.
    pub fn new(name: &'a str) -> Self {
        Command {
            hidden: false,
            priority: PRIORITY_DEFAULT,
            name: name,
            help_text: None,
            parameters: vec![],
            wrapped_root: None,
        }
    }

    /// Mark the command as hidden. Hidden commands will match
    /// within the parser, but are not listed during completion.
    pub fn hidden(mut self, hidden: bool) -> Self {
        self.hidden = hidden;
        self
    }

    /// Give the command a priority. This is used when sorting
    /// out conflicts during matching and completion.
    pub fn priority(mut self, priority: i32) -> Self {
        self.priority = priority;
        self
    }

    /// Supply help text for the command.
    pub fn help(mut self, help_text: &'a str) -> Self {
        self.help_text = Some(help_text);
        self
    }

    /// Add a [`Parameter`] to the command.
    ///
    /// [`Parameter`]: struct.Parameter.html
    pub fn parameter(mut self, parameter: Parameter<'a>) -> Self {
        self.parameters.push(parameter);
        self
    }

    /// The `wrapped_root` signifies the path to the command that should
    /// be wrapped by this command. This is used for the `help` command.
    ///
    /// [`CommandNode`]: struct.CommandNode.html
    pub fn wraps(mut self, wrapped_root: String) -> Self {
        self.wrapped_root = Some(wrapped_root);
        self
    }
}

/// Description of a parameter to be added to the [`Command`].
///
/// The lifetime parameter `'a` refers to the lifetime
/// of the strings used for parameter names, aliases and
/// help text.
///
/// [`Command`]: struct.Command.html
pub struct Parameter<'a> {
    hidden: bool,
    priority: Option<i32>,
    name: &'a str,
    repeatable: bool,
    aliases: Vec<&'a str>,
    help_text: Option<&'a str>,
    kind: ParameterKind,
    required: bool,
}

impl<'a> Parameter<'a> {
    /// Construct a default (blank) parameter with the given `name`.
    pub fn new(name: &'a str) -> Self {
        Parameter {
            hidden: false,
            priority: None,
            name: name,
            repeatable: false,
            aliases: vec![],
            help_text: None,
            kind: ParameterKind::Simple,
            required: false,
        }
    }

    /// Mark the parameter as hidden. Hidden parameters will match
    /// within the parser, but are not listed during completion.
    pub fn hidden(mut self, hidden: bool) -> Self {
        self.hidden = hidden;
        self
    }

    /// Give the parameter a priority. This is used when sorting
    /// out conflicts during matching and completion.
    ///
    /// The `priority` of a `Parameter` defaults to `PRIORITY_PARAMETER`
    /// except for when the `kind` is `ParameterKind::Flag` in which
    /// case, the default will be `PRIORITY_DEFAULT`.
    pub fn priority(mut self, priority: i32) -> Self {
        self.priority = Some(priority);
        self
    }

    /// Establish whether or not this parameter is repeatable.
    /// Repeated parameters produce a vector of values and can
    /// be given multiple times within a single command invocation.
    pub fn repeatable(mut self, repeatable: bool) -> Self {
        self.repeatable = repeatable;
        self
    }

    /// Add an alias that this parameter can use.
    ///
    /// Aliases are currently only valid for parameters of `kind`
    /// `ParameterKind::Named`.
    pub fn alias(mut self, alias: &'a str) -> Self {
        self.aliases.push(alias);
        self
    }

    /// Supply the help text for the parameter.
    pub fn help(mut self, help_text: &'a str) -> Self {
        self.help_text = Some(help_text);
        self
    }

    /// Establish whether or not this parameter is required.
    pub fn required(mut self, required: bool) -> Self {
        self.required = required;
        self
    }

    /// Set which type of [`ParameterNode`] is supposed to be created
    /// to represent this parameter.
    ///
    /// [`ParameterNode`]: trait.ParameterNode.html
    pub fn kind(mut self, kind: ParameterKind) -> Self {
        self.kind = kind;
        self
    }
}
<|endoftext|>