<commit_before>use nom;
use nom::{Err, IResult};

use metadata;
use frame;
use subframe;

use metadata::{Metadata, StreamInfo, metadata_parser};
use frame::{frame_parser, Frame};
use utility::{ErrorKind, ByteStream, ReadStream, StreamProducer};

use std::io;
use std::usize;
use std::fs::File;

enum ParserState {
  Marker,
  Metadata,
  Frame,
}

/// FLAC stream that decodes and hold file information.
pub struct Stream {
  info: StreamInfo,
  metadata: Vec<Metadata>,
  frames: Vec<Frame>,
  state: ParserState,
  output: Vec<i32>,
  frame_index: usize,
}

named!(pub stream_parser <&[u8], Stream>,
  chain!(
    blocks: metadata_parser ~
    frames: many1!(apply!(frame_parser, &blocks.0)),
    move|| {
      Stream {
        info: blocks.0,
        metadata: blocks.1,
        frames: frames,
        state: ParserState::Marker,
        output: Vec::new(),
        frame_index: 0,
      }
    }
  )
);

impl Stream {
  pub fn new() -> Stream {
    Stream {
      info: StreamInfo::new(),
      metadata: Vec::new(),
      frames: Vec::new(),
      state: ParserState::Marker,
      output: Vec::new(),
      frame_index: 0,
    }
  }

  pub fn info(&self) -> StreamInfo {
    self.info
  }

  pub fn metadata(&self) -> &[Metadata] {
    &self.metadata
  }

  pub fn from_file(filename: &str) -> io::Result<Stream> {
    File::open(filename).and_then(|file| {
      let mut producer = ReadStream::new(file);
      let error_str    = format!("parser: couldn't parse the given file {}",
                                 filename);

      Stream::from_stream_producer(&mut producer, &error_str)
    })
  }

  pub fn from_buffer(buffer: &[u8]) -> io::Result<Stream> {
    let mut producer = ByteStream::new(buffer);
    let error_str    = "parser: couldn't parse the buffer";

    Stream::from_stream_producer(&mut producer, error_str)
  }

  fn from_stream_producer<P>(producer: &mut P, error_str: &str)
                             -> io::Result<Stream>
   where P: StreamProducer {
    let mut is_error = false;
    let mut stream   = Stream {
      info: StreamInfo::new(),
      metadata: Vec::new(),
      frames: Vec::new(),
      state: ParserState::Marker,
      output: Vec::new(),
      frame_index: 0,
    };

    loop {
      match stream.handle(producer) {
        Ok(_)                         => break,
        Err(ErrorKind::EndOfInput)    => break,
        Err(ErrorKind::Consumed(_))   => continue,
        Err(ErrorKind::Incomplete(_)) => continue,
        Err(_)                        => {
          is_error = true;

          break;
        }
      }
    }

    if !is_error {
      let channels    = stream.info.channels as usize;
      let block_size  = stream.info.max_block_size as usize;
      let output_size = block_size * channels;

      stream.output.reserve_exact(output_size);

      unsafe { stream.output.set_len(output_size) }

      Ok(stream)
    } else {
      Err(io::Error::new(io::ErrorKind::InvalidData, error_str))
    }
  }

  pub fn iter(&mut self) -> Iter {
    Iter::new(self)
  }

  fn next_frame<'a>(&'a mut self) -> Option<&'a [i32]> {
    if self.frames.is_empty() || self.frame_index >= self.frames.len() {
      None
    } else {
      let frame       = &self.frames[self.frame_index];
      let channels    = frame.header.channels as usize;
      let block_size  = frame.header.block_size as usize;
      let mut channel = 0;

      for subframe in &frame.subframes[0..channels] {
        let start  = channel * block_size;
        let end    = (channel + 1) * block_size;
        let output = &mut self.output[start..end];

        subframe::decode(&subframe, output);

        channel += 1;
      }

      frame::decode(frame.header.channel_assignment, &mut self.output);

      self.frame_index += 1;

      Some(&self.output[0..(block_size * channels)])
    }
  }

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

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

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

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

    match metadata::block(input) {
      IResult::Done(i, block) => {
        let is_last = block.is_last;

        if let metadata::Data::StreamInfo(info) = block.data {
          self.info = info;
        } else {
          self.metadata.push(block);
        }

        if is_last {
          self.state = ParserState::Frame;
        }

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

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

    match frame_parser(input, &self.info) {
      IResult::Done(i, frame) => {
        self.frames.push(frame);

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

  fn handle<S: StreamProducer>(&mut self, stream: &mut S)
                               -> Result<(), ErrorKind> {
    stream.parse(|input| {
      match self.state {
        ParserState::Marker   => self.handle_marker(input),
        ParserState::Metadata => self.handle_metadata(input),
        ParserState::Frame    => self.handle_frame(input),
      }
    })
  }
}

pub struct Iter<'a> {
  stream: &'a mut Stream,
  channel: usize,
  frame_index: usize,
  block_size: usize,
  sample_index: usize,
  samples_left: u64,
}

impl<'a> Iter<'a> {
  pub fn new(stream: &'a mut Stream) -> Iter<'a> {
    let samples_left = stream.info.total_samples;

    Iter {
      stream: stream,
      channel: 0,
      frame_index: 0,
      block_size: 0,
      sample_index: 0,
      samples_left: samples_left,
    }
  }
}

impl<'a> Iterator for Iter<'a> {
  type Item = i32;

  fn next(&mut self) -> Option<Self::Item> {
    if self.block_size == 0 || self.sample_index == self.block_size {
      if self.stream.next_frame().is_none() {
        return None;
      } else {
        let frame = &self.stream.frames[self.frame_index];

        self.sample_index = 0;
        self.block_size   = frame.header.block_size as usize;
      }
    }

    let channels = self.stream.info.channels as usize;
    let index    = self.sample_index + (self.channel * self.block_size);
    let sample   = self.stream.output[index];

    self.channel      += 1;
    self.samples_left -= 1;

    // Reset current channel
    if self.channel == channels {
      self.channel       = 0;
      self.sample_index += 1;
    }

    Some(sample)
  }

  fn size_hint(&self) -> (usize, Option<usize>) {
    let samples_left = self.samples_left as usize;
    let max_value    = usize::max_value() as u64;

    // There is a chance that samples_left will be larger than a usize since
    // it is a u64. Make the upper bound None when it is.
    if self.samples_left > max_value {
      (samples_left, None)
    } else {
      (samples_left, Some(samples_left))
    }
  }
}

impl<'a> IntoIterator for &'a mut Stream {
  type Item     = i32;
  type IntoIter = Iter<'a>;

  fn into_iter(self) -> Self::IntoIter {
    self.iter()
  }
}
<commit_msg>Add documentation for all three constructors<commit_after>use nom;
use nom::{Err, IResult};

use metadata;
use frame;
use subframe;

use metadata::{Metadata, StreamInfo, metadata_parser};
use frame::{frame_parser, Frame};
use utility::{ErrorKind, ByteStream, ReadStream, StreamProducer};

use std::io;
use std::usize;
use std::fs::File;

enum ParserState {
  Marker,
  Metadata,
  Frame,
}

/// FLAC stream that decodes and hold file information.
pub struct Stream {
  info: StreamInfo,
  metadata: Vec<Metadata>,
  frames: Vec<Frame>,
  state: ParserState,
  output: Vec<i32>,
  frame_index: usize,
}

named!(pub stream_parser <&[u8], Stream>,
  chain!(
    blocks: metadata_parser ~
    frames: many1!(apply!(frame_parser, &blocks.0)),
    move|| {
      Stream {
        info: blocks.0,
        metadata: blocks.1,
        frames: frames,
        state: ParserState::Marker,
        output: Vec::new(),
        frame_index: 0,
      }
    }
  )
);

impl Stream {
  /// Constructor for the default state of a FLAC stream.
  ///
  /// This doesn't actually decode anything, it just hold the default values
  /// of each field.
  pub fn new() -> Stream {
    Stream {
      info: StreamInfo::new(),
      metadata: Vec::new(),
      frames: Vec::new(),
      state: ParserState::Marker,
      output: Vec::new(),
      frame_index: 0,
    }
  }

  pub fn info(&self) -> StreamInfo {
    self.info
  }

  pub fn metadata(&self) -> &[Metadata] {
    &self.metadata
  }

  /// Constructs a decoder with the given file name.
  ///
  /// # Failures
  ///
  /// * `ErrorKind::NotFound` is returned when the given filename isn't found.
  /// * `ErrorKind::InvalidData` is returned when the data within the file
  ///   isn't valid FLAC data.
  pub fn from_file(filename: &str) -> io::Result<Stream> {
    File::open(filename).and_then(|file| {
      let mut producer = ReadStream::new(file);
      let error_str    = format!("parser: couldn't parse the given file {}",
                                 filename);

      Stream::from_stream_producer(&mut producer, &error_str)
    })
  }

  /// Constructs a decoder with the given buffer.
  ///
  /// This constructor assumes that an entire FLAC file is in the buffer.
  ///
  /// # Failures
  ///
  /// * `ErrorKind::InvalidData` is returned when the data within the buffer
  ///   isn't valid FLAC data.
  pub fn from_buffer(buffer: &[u8]) -> io::Result<Stream> {
    let mut producer = ByteStream::new(buffer);
    let error_str    = "parser: couldn't parse the buffer";

    Stream::from_stream_producer(&mut producer, error_str)
  }

  fn from_stream_producer<P>(producer: &mut P, error_str: &str)
                             -> io::Result<Stream>
   where P: StreamProducer {
    let mut is_error = false;
    let mut stream   = Stream {
      info: StreamInfo::new(),
      metadata: Vec::new(),
      frames: Vec::new(),
      state: ParserState::Marker,
      output: Vec::new(),
      frame_index: 0,
    };

    loop {
      match stream.handle(producer) {
        Ok(_)                         => break,
        Err(ErrorKind::EndOfInput)    => break,
        Err(ErrorKind::Consumed(_))   => continue,
        Err(ErrorKind::Incomplete(_)) => continue,
        Err(_)                        => {
          is_error = true;

          break;
        }
      }
    }

    if !is_error {
      let channels    = stream.info.channels as usize;
      let block_size  = stream.info.max_block_size as usize;
      let output_size = block_size * channels;

      stream.output.reserve_exact(output_size);

      unsafe { stream.output.set_len(output_size) }

      Ok(stream)
    } else {
      Err(io::Error::new(io::ErrorKind::InvalidData, error_str))
    }
  }

  pub fn iter(&mut self) -> Iter {
    Iter::new(self)
  }

  fn next_frame<'a>(&'a mut self) -> Option<&'a [i32]> {
    if self.frames.is_empty() || self.frame_index >= self.frames.len() {
      None
    } else {
      let frame       = &self.frames[self.frame_index];
      let channels    = frame.header.channels as usize;
      let block_size  = frame.header.block_size as usize;
      let mut channel = 0;

      for subframe in &frame.subframes[0..channels] {
        let start  = channel * block_size;
        let end    = (channel + 1) * block_size;
        let output = &mut self.output[start..end];

        subframe::decode(&subframe, output);

        channel += 1;
      }

      frame::decode(frame.header.channel_assignment, &mut self.output);

      self.frame_index += 1;

      Some(&self.output[0..(block_size * channels)])
    }
  }

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

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

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

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

    match metadata::block(input) {
      IResult::Done(i, block) => {
        let is_last = block.is_last;

        if let metadata::Data::StreamInfo(info) = block.data {
          self.info = info;
        } else {
          self.metadata.push(block);
        }

        if is_last {
          self.state = ParserState::Frame;
        }

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

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

    match frame_parser(input, &self.info) {
      IResult::Done(i, frame) => {
        self.frames.push(frame);

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

  fn handle<S: StreamProducer>(&mut self, stream: &mut S)
                               -> Result<(), ErrorKind> {
    stream.parse(|input| {
      match self.state {
        ParserState::Marker   => self.handle_marker(input),
        ParserState::Metadata => self.handle_metadata(input),
        ParserState::Frame    => self.handle_frame(input),
      }
    })
  }
}

pub struct Iter<'a> {
  stream: &'a mut Stream,
  channel: usize,
  frame_index: usize,
  block_size: usize,
  sample_index: usize,
  samples_left: u64,
}

impl<'a> Iter<'a> {
  pub fn new(stream: &'a mut Stream) -> Iter<'a> {
    let samples_left = stream.info.total_samples;

    Iter {
      stream: stream,
      channel: 0,
      frame_index: 0,
      block_size: 0,
      sample_index: 0,
      samples_left: samples_left,
    }
  }
}

impl<'a> Iterator for Iter<'a> {
  type Item = i32;

  fn next(&mut self) -> Option<Self::Item> {
    if self.block_size == 0 || self.sample_index == self.block_size {
      if self.stream.next_frame().is_none() {
        return None;
      } else {
        let frame = &self.stream.frames[self.frame_index];

        self.sample_index = 0;
        self.block_size   = frame.header.block_size as usize;
      }
    }

    let channels = self.stream.info.channels as usize;
    let index    = self.sample_index + (self.channel * self.block_size);
    let sample   = self.stream.output[index];

    self.channel      += 1;
    self.samples_left -= 1;

    // Reset current channel
    if self.channel == channels {
      self.channel       = 0;
      self.sample_index += 1;
    }

    Some(sample)
  }

  fn size_hint(&self) -> (usize, Option<usize>) {
    let samples_left = self.samples_left as usize;
    let max_value    = usize::max_value() as u64;

    // There is a chance that samples_left will be larger than a usize since
    // it is a u64. Make the upper bound None when it is.
    if self.samples_left > max_value {
      (samples_left, None)
    } else {
      (samples_left, Some(samples_left))
    }
  }
}

impl<'a> IntoIterator for &'a mut Stream {
  type Item     = i32;
  type IntoIter = Iter<'a>;

  fn into_iter(self) -> Self::IntoIter {
    self.iter()
  }
}
<|endoftext|>


<commit_before>use nom;
use nom::{Err, IResult};

use metadata;
use frame;
use subframe;

use metadata::{Metadata, StreamInfo, metadata_parser};
use frame::{frame_parser, Frame};
use utility::{ErrorKind, ByteStream, ReadStream, StreamProducer};

use std::io;
use std::usize;
use std::fs::File;

enum ParserState {
  Marker,
  Metadata,
  Frame,
}

/// FLAC stream that decodes and hold file information.
pub struct Stream
 {
  info: StreamInfo,
  metadata: Vec<Metadata>,
  frames: Vec<Frame>,
  state: ParserState,
  output: Vec<i32>,
  frame_index: usize,
  producer: P,
}

fn parser<'a>(input: &'a [u8], is_start: &mut bool)
              -> IResult<&'a [u8], Metadata> {
  let mut slice = input;

  if *is_start {
    let (i, _) = try_parse!(slice, tag!("fLaC"));

    slice     = i;
    *is_start = false;
  }

  metadata_parser(slice)
}

impl
 Stream
 where P: StreamProducer {
  /// Constructor for the default state of a FLAC stream.
  ///
  /// This doesn't actually decode anything, it just hold the default values
  /// of each field.
  pub fn new() -> Stream {
    Stream {
      info: StreamInfo::new(),
      metadata: Vec::new(),
      frames: Vec::new(),
      state: ParserState::Marker,
      output: Vec::new(),
      frame_index: 0,
    }
  }

  /// Returns information for the current stream.
  pub fn info(&self) -> StreamInfo {
    self.info
  }

  /// Returns a slice of `Metadata`
  ///
  /// This slice excludes `StreamInfo`, which is located in `Stream::info`.
  /// Everything else is related to metadata for the FLAC stream is in the
  /// slice.
  pub fn metadata(&self) -> &[Metadata] {
    &self.metadata
  }

  /// Constructs a decoder with the given file name.
  ///
  /// # Failures
  ///
  /// * `ErrorKind::NotFound` is returned when the given filename isn't found.
  /// * `ErrorKind::InvalidData` is returned when the data within the file
  ///   isn't valid FLAC data.
  pub fn from_file(filename: &str) -> io::Result<Stream<ReadStream<File>>> {
    File::open(filename).and_then(|file| {
      let producer  = ReadStream::new(file);
      let error_str = format!("parser: couldn't parse the given file {}",
                              filename);

      Stream::from_stream_producer(producer, &error_str)
    })
  }

  /// Constructs a decoder with the given buffer.
  ///
  /// This constructor assumes that an entire FLAC file is in the buffer.
  ///
  /// # Failures
  ///
  /// * `ErrorKind::InvalidData` is returned when the data within the buffer
  ///   isn't valid FLAC data.
  pub fn from_buffer(buffer: &[u8]) -> io::Result<Stream<ByteStream>> {
    let producer  = ByteStream::new(buffer);
    let error_str = "parser: couldn't parse the buffer";

    Stream::from_stream_producer(producer, error_str)
  }

  fn from_stream_producer(mut producer: P, error_str: &str)
                          -> io::Result<Self> {
    let mut is_start    = true;
    let mut is_error    = false;
    let mut stream_info = StreamInfo::new();
    let mut metadata    = Vec::new();

    loop {
      match producer.parse(|i| parser(i, &mut is_start)) {
        Ok(block)                  => {
          let is_last = block.is_last;

          if let metadata::Data::StreamInfo(info) = block.data {
            stream_info = info;
          } else {
            metadata.push(block);
          }

          if is_last {
            break;
          }
        }
        Err(ErrorKind::EndOfInput) => break,
        Err(ErrorKind::Continue)   => continue,
        Err(_)                     => {
          is_error = true;

          break;
        }
      }
    }

    if !is_error {
      let channels    = stream_info.channels as usize;
      let block_size  = stream_info.max_block_size as usize;
      let output_size = block_size * channels;
      let mut output  = Vec::with_capacity(output_size);

      unsafe { output.set_len(output_size) }

      Ok(Stream {
        info: stream_info,
        metadata: metadata,
        frames: Vec::new(),
        state: ParserState::Marker,
        output: output,
        frame_index: 0,
        producer: producer,
      })
    } else {
      Err(io::Error::new(io::ErrorKind::InvalidData, error_str))
    }
  }

  /// Returns an iterator over the decoded samples.
  pub fn iter(&mut self) -> Iter
 {
    let samples_left = self.info.total_samples;

    Iter {
      stream: self,
      channel: 0,
      block_size: 0,
      sample_index: 0,
      samples_left: samples_left,
    }
  }

  fn next_frame<'a>(&'a mut self) -> Option<usize> {
    let stream_info = self.info();

    loop {
      match self.producer.parse(|i| frame_parser(i, &stream_info)) {
        Ok(frame)                  => {
          let channels    = frame.header.channels as usize;
          let block_size  = frame.header.block_size as usize;
          let mut channel = 0;

          for subframe in &frame.subframes[0..channels] {
            let start  = channel * block_size;
            let end    = (channel + 1) * block_size;
            let output = &mut self.output[start..end];

            subframe::decode(&subframe, output);

            channel += 1;
          }

          frame::decode(frame.header.channel_assignment, &mut self.output);

          self.frame_index += 1;

          return Some(block_size);
        }
        Err(ErrorKind::EndOfInput) => return None,
        Err(ErrorKind::Continue)   => continue,
        Err(_)                     => return None,
      }
    }
  }

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

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

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

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

    match metadata_parser(input) {
      IResult::Done(i, block) => {
        let is_last = block.is_last;

        if let metadata::Data::StreamInfo(info) = block.data {
          self.info = info;
        } else {
          self.metadata.push(block);
        }

        if is_last {
          self.state = ParserState::Frame;
        }

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

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

    match frame_parser(input, &self.info) {
      IResult::Done(i, frame) => {
        self.frames.push(frame);

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

  fn handle<S: StreamProducer>(&mut self, stream: &mut S)
                               -> Result<(), ErrorKind> {
    stream.parse(|input| {
      match self.state {
        ParserState::Marker   => self.handle_marker(input),
        ParserState::Metadata => self.handle_metadata(input),
        ParserState::Frame    => self.handle_frame(input),
      }
    })
  }
}

/// An iterator over a reference of the decoded FLAC stream.
pub struct Iter<'a, P> where P: 'a + StreamProducer {
  stream: &'a mut Stream
,
  channel: usize,
  block_size: usize,
  sample_index: usize,
  samples_left: u64,
}

impl<'a, P> Iterator for Iter<'a, P> where P: StreamProducer {
  type Item = i32;

  fn next(&mut self) -> Option<Self::Item> {
    if self.sample_index == self.block_size {
      if let Some(block_size) = self.stream.next_frame() {
        self.sample_index = 0;
        self.block_size   = block_size;
      } else {
        return None;
      }
    }

    let channels = self.stream.info.channels as usize;
    let index    = self.sample_index + (self.channel * self.block_size);
    let sample   = self.stream.output[index];

    self.channel += 1;

    // Reset current channel
    if self.channel == channels {
      self.channel       = 0;
      self.sample_index += 1;
      self.samples_left -= 1;
    }

    Some(sample)
  }

  fn size_hint(&self) -> (usize, Option<usize>) {
    let samples_left = self.samples_left as usize;
    let max_value    = usize::max_value() as u64;

    // There is a chance that samples_left will be larger than a usize since
    // it is a u64. Make the upper bound None when it is.
    if self.samples_left > max_value {
      (samples_left, None)
    } else {
      (samples_left, Some(samples_left))
    }
  }
}

impl<'a, P> IntoIterator for &'a mut Stream

 where P: StreamProducer {
  type Item     = i32;
  type IntoIter = Iter<'a, P>;

  fn into_iter(self) -> Self::IntoIter {
    self.iter()
  }
}
<commit_msg>Rewrite `Stream::new` to use the `Read` trait<commit_after>use nom;
use nom::{Err, IResult};

use metadata;
use frame;
use subframe;

use metadata::{Metadata, StreamInfo, metadata_parser};
use frame::{frame_parser, Frame};
use utility::{ErrorKind, ByteStream, ReadStream, StreamProducer};

use std::io;
use std::usize;
use std::fs::File;

enum ParserState {
  Marker,
  Metadata,
  Frame,
}

/// FLAC stream that decodes and hold file information.
pub struct Stream
 {
  info: StreamInfo,
  metadata: Vec<Metadata>,
  frames: Vec<Frame>,
  state: ParserState,
  output: Vec<i32>,
  frame_index: usize,
  producer: P,
}

fn parser<'a>(input: &'a [u8], is_start: &mut bool)
              -> IResult<&'a [u8], Metadata> {
  let mut slice = input;

  if *is_start {
    let (i, _) = try_parse!(slice, tag!("fLaC"));

    slice     = i;
    *is_start = false;
  }

  metadata_parser(slice)
}

impl
 Stream
 where P: StreamProducer {
  /// Constructor for the default state of a FLAC stream.
  ///
  /// This doesn't actually decode anything, it just hold the default values
  /// of each field.
  pub fn new<R: io::Read>(reader: R) -> io::Result<Stream<ReadStream<R>>> {
    let producer  = ReadStream::new(reader);
    let error_str = "parser: couldn't parse the reader";

    Stream::from_stream_producer(producer, error_str)
  }

  /// Returns information for the current stream.
  pub fn info(&self) -> StreamInfo {
    self.info
  }

  /// Returns a slice of `Metadata`
  ///
  /// This slice excludes `StreamInfo`, which is located in `Stream::info`.
  /// Everything else is related to metadata for the FLAC stream is in the
  /// slice.
  pub fn metadata(&self) -> &[Metadata] {
    &self.metadata
  }

  /// Constructs a decoder with the given file name.
  ///
  /// # Failures
  ///
  /// * `ErrorKind::NotFound` is returned when the given filename isn't found.
  /// * `ErrorKind::InvalidData` is returned when the data within the file
  ///   isn't valid FLAC data.
  pub fn from_file(filename: &str) -> io::Result<Stream<ReadStream<File>>> {
    File::open(filename).and_then(|file| {
      let producer  = ReadStream::new(file);
      let error_str = format!("parser: couldn't parse the given file {}",
                              filename);

      Stream::from_stream_producer(producer, &error_str)
    })
  }

  /// Constructs a decoder with the given buffer.
  ///
  /// This constructor assumes that an entire FLAC file is in the buffer.
  ///
  /// # Failures
  ///
  /// * `ErrorKind::InvalidData` is returned when the data within the buffer
  ///   isn't valid FLAC data.
  pub fn from_buffer(buffer: &[u8]) -> io::Result<Stream<ByteStream>> {
    let producer  = ByteStream::new(buffer);
    let error_str = "parser: couldn't parse the buffer";

    Stream::from_stream_producer(producer, error_str)
  }

  fn from_stream_producer(mut producer: P, error_str: &str)
                          -> io::Result<Self> {
    let mut is_start    = true;
    let mut is_error    = false;
    let mut stream_info = StreamInfo::new();
    let mut metadata    = Vec::new();

    loop {
      match producer.parse(|i| parser(i, &mut is_start)) {
        Ok(block)                  => {
          let is_last = block.is_last;

          if let metadata::Data::StreamInfo(info) = block.data {
            stream_info = info;
          } else {
            metadata.push(block);
          }

          if is_last {
            break;
          }
        }
        Err(ErrorKind::EndOfInput) => break,
        Err(ErrorKind::Continue)   => continue,
        Err(_)                     => {
          is_error = true;

          break;
        }
      }
    }

    if !is_error {
      let channels    = stream_info.channels as usize;
      let block_size  = stream_info.max_block_size as usize;
      let output_size = block_size * channels;
      let mut output  = Vec::with_capacity(output_size);

      unsafe { output.set_len(output_size) }

      Ok(Stream {
        info: stream_info,
        metadata: metadata,
        frames: Vec::new(),
        state: ParserState::Marker,
        output: output,
        frame_index: 0,
        producer: producer,
      })
    } else {
      Err(io::Error::new(io::ErrorKind::InvalidData, error_str))
    }
  }

  /// Returns an iterator over the decoded samples.
  pub fn iter(&mut self) -> Iter
 {
    let samples_left = self.info.total_samples;

    Iter {
      stream: self,
      channel: 0,
      block_size: 0,
      sample_index: 0,
      samples_left: samples_left,
    }
  }

  fn next_frame<'a>(&'a mut self) -> Option<usize> {
    let stream_info = self.info();

    loop {
      match self.producer.parse(|i| frame_parser(i, &stream_info)) {
        Ok(frame)                  => {
          let channels    = frame.header.channels as usize;
          let block_size  = frame.header.block_size as usize;
          let mut channel = 0;

          for subframe in &frame.subframes[0..channels] {
            let start  = channel * block_size;
            let end    = (channel + 1) * block_size;
            let output = &mut self.output[start..end];

            subframe::decode(&subframe, output);

            channel += 1;
          }

          frame::decode(frame.header.channel_assignment, &mut self.output);

          self.frame_index += 1;

          return Some(block_size);
        }
        Err(ErrorKind::EndOfInput) => return None,
        Err(ErrorKind::Continue)   => continue,
        Err(_)                     => return None,
      }
    }
  }

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

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

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

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

    match metadata_parser(input) {
      IResult::Done(i, block) => {
        let is_last = block.is_last;

        if let metadata::Data::StreamInfo(info) = block.data {
          self.info = info;
        } else {
          self.metadata.push(block);
        }

        if is_last {
          self.state = ParserState::Frame;
        }

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

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

    match frame_parser(input, &self.info) {
      IResult::Done(i, frame) => {
        self.frames.push(frame);

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

  fn handle<S: StreamProducer>(&mut self, stream: &mut S)
                               -> Result<(), ErrorKind> {
    stream.parse(|input| {
      match self.state {
        ParserState::Marker   => self.handle_marker(input),
        ParserState::Metadata => self.handle_metadata(input),
        ParserState::Frame    => self.handle_frame(input),
      }
    })
  }
}

/// An iterator over a reference of the decoded FLAC stream.
pub struct Iter<'a, P> where P: 'a + StreamProducer {
  stream: &'a mut Stream
,
  channel: usize,
  block_size: usize,
  sample_index: usize,
  samples_left: u64,
}

impl<'a, P> Iterator for Iter<'a, P> where P: StreamProducer {
  type Item = i32;

  fn next(&mut self) -> Option<Self::Item> {
    if self.sample_index == self.block_size {
      if let Some(block_size) = self.stream.next_frame() {
        self.sample_index = 0;
        self.block_size   = block_size;
      } else {
        return None;
      }
    }

    let channels = self.stream.info.channels as usize;
    let index    = self.sample_index + (self.channel * self.block_size);
    let sample   = self.stream.output[index];

    self.channel += 1;

    // Reset current channel
    if self.channel == channels {
      self.channel       = 0;
      self.sample_index += 1;
      self.samples_left -= 1;
    }

    Some(sample)
  }

  fn size_hint(&self) -> (usize, Option<usize>) {
    let samples_left = self.samples_left as usize;
    let max_value    = usize::max_value() as u64;

    // There is a chance that samples_left will be larger than a usize since
    // it is a u64. Make the upper bound None when it is.
    if self.samples_left > max_value {
      (samples_left, None)
    } else {
      (samples_left, Some(samples_left))
    }
  }
}

impl<'a, P> IntoIterator for &'a mut Stream

 where P: StreamProducer {
  type Item     = i32;
  type IntoIter = Iter<'a, P>;

  fn into_iter(self) -> Self::IntoIter {
    self.iter()
  }
}
<|endoftext|>


<commit_before>use std::env;
use std::fs::File;
use std::io::{Read, Write};

fn main() {
    let mut bs = 512;
    let mut count = None;
    let mut in_path = String::new();
    let mut out_path = String::new();
    for arg in env::args().skip(1) {
        if arg.starts_with("bs=") {
            bs = arg[3..].parse::<usize>().expect("dd: bs is not a number");
        } else if arg.starts_with("count=") {
            count = Some(arg[6..].parse::<usize>().expect("dd: count is not a number"));
        } else if arg.starts_with("if=") {
            in_path = arg[3..].to_string();
        } else if arg.starts_with("of=") {
            out_path = arg[3..].to_string();
        } else {
            panic!("dd: unrecognized operand '{}'", arg);
        }
    }

    let mut input = File::open(in_path).expect("dd: failed to open if");
    let mut output = File::create(out_path).expect("dd: failed to open of");

    let mut in_recs = 0;
    let mut in_extra = 0;
    let mut out_recs = 0;
    let mut out_extra = 0;
    let mut out_total = 0;
    let mut running = true;
    let mut buffer = vec![0; bs];
    while running {
        let in_count = input.read(&mut buffer).expect("dd: failed to read if");
        if in_count < bs {
            if in_count > 0 {
                in_extra += 1;
            }
            running = false;
        } else {
            in_recs += 1;
            if let Some(count) = count {
                if in_recs >= count {
                    running = false;
                }
            }
        }

        let out_count = output.write(&buffer[.. in_count]).expect("dd: failed to write of");
        if out_count < bs {
            if out_count > 0 {
                out_extra += 1;
            }
            running = false;
        } else {
            out_recs += 1;
        }

        out_total += out_count;
    }

    println!("{}+{} records in", in_recs, in_extra);
    println!("{}+{} records out", out_recs, out_extra);
    println!("{} bytes copied", out_total);
}
<commit_msg>Better information from dd<commit_after>extern crate coreutils;

use coreutils::to_human_readable_string;
use std::env;
use std::fs::File;
use std::io::{stderr, Read, Write};
use std::time::Instant;

fn main() {
    let stderr = stderr();
    let mut stderr = stderr.lock();

    let mut bs = 512;
    let mut count = None;
    let mut in_path = String::new();
    let mut out_path = String::new();
    let mut status = 1;
    for arg in env::args().skip(1) {
        if arg.starts_with("bs=") {
            bs = arg[3..].parse::<usize>().expect("dd: bs is not a number");
        } else if arg.starts_with("count=") {
            count = Some(arg[6..].parse::<usize>().expect("dd: count is not a number"));
        } else if arg.starts_with("if=") {
            in_path = arg[3..].to_string();
        } else if arg.starts_with("of=") {
            out_path = arg[3..].to_string();
        } else if arg.starts_with("status=") {
            match &arg[7..] {
                "none" => status = 0,
                "noxfer" => status = 1,
                "progress" => status = 2,
                unknown => panic!("dd: status: unrecognized argument '{}'", unknown)
            }
        } else {
            panic!("dd: unrecognized operand '{}'", arg);
        }
    }

    let mut input = File::open(in_path).expect("dd: failed to open if");
    let mut output = File::create(out_path).expect("dd: failed to open of");

    let mut in_recs = 0;
    let mut in_extra = 0;
    let mut out_recs = 0;
    let mut out_extra = 0;
    let mut out_total = 0;
    let mut running = true;
    let mut buffer = vec![0; bs];
    let mut start = Instant::now();
    let mut last_print = 0;
    let mut last_print_out = 0;
    while running {
        let in_count = input.read(&mut buffer).expect("dd: failed to read if");
        if in_count < bs {
            if in_count > 0 {
                in_extra += 1;
            }
            running = false;
        } else {
            in_recs += 1;
            if let Some(count) = count {
                if in_recs >= count {
                    running = false;
                }
            }
        }

        let out_count = output.write(&buffer[.. in_count]).expect("dd: failed to write of");
        if out_count < bs {
            if out_count > 0 {
                out_extra += 1;
            }
            running = false;
        } else {
            out_recs += 1;
        }

        out_total += out_count as u64;

        if status >= 2 {
            let elapsed = start.elapsed().as_secs();
            if elapsed > last_print {
                let _ = write!(stderr, "\r\x1B[K{} bytes ({}B) copied, {} s, {}B/s", out_total, to_human_readable_string(out_total), elapsed, to_human_readable_string(out_total - last_print_out));
                let _ = stderr.flush();
                last_print = elapsed;
                last_print_out = out_total;
            }
        }
    }

    if status >= 1 {
        let elapsed_duration = start.elapsed();
        let elapsed = elapsed_duration.as_secs() as f64 + (elapsed_duration.subsec_nanos() as f64)/1000000000.0;

        if status >= 2 && last_print > 0 {
            let _ = writeln!(stderr, "");
        }

        let _ = writeln!(stderr, "{}+{} records in", in_recs, in_extra);
        let _ = writeln!(stderr, "{}+{} records out", out_recs, out_extra);
        let _ = writeln!(stderr, "{} bytes ({}B) copied, {} s, {}B/s", out_total, to_human_readable_string(out_total), elapsed, to_human_readable_string((out_total as f64/elapsed) as u64));
    }
}
<|endoftext|>


<commit_before>// compile-flags: -C no-prepopulate-passes
// ignore-tidy-linelength

// min-system-llvm-version: 9.0
// ignore-arm
// ignore-mips
// ignore-mips64
// ignore-powerpc
// ignore-powerpc64
// ignore-riscv64 see codegen/riscv-abi
// ignore-windows
// See repr-transparent.rs

#![feature(transparent_unions)]

#![crate_type="lib"]

#[derive(Clone, Copy)]
#[repr(C)]
pub struct BigS([u32; 16]);

#[repr(transparent)]
pub struct TsBigS(BigS);

#[repr(transparent)]
pub union TuBigS {
    field: BigS,
}

#[repr(transparent)]
pub enum TeBigS {
    Variant(BigS),
}

// CHECK: define void @test_BigS(%BigS* [[BIGS_RET_ATTRS:.*]], %BigS* [[BIGS_ARG_ATTRS1:.*]] byval(%BigS) [[BIGS_ARG_ATTRS2:.*]])
#[no_mangle]
pub extern fn test_BigS(_: BigS) -> BigS { loop {} }

// CHECK: define void @test_TsBigS(%TsBigS* [[BIGS_RET_ATTRS]], %TsBigS* [[BIGS_ARG_ATTRS1]] byval(%TsBigS) [[BIGS_ARG_ATTRS2:.*]])
#[no_mangle]
pub extern fn test_TsBigS(_: TsBigS) -> TsBigS { loop {} }

// CHECK: define void @test_TuBigS(%TuBigS* [[BIGS_RET_ATTRS]], %TuBigS* [[BIGS_ARG_ATTRS1]] byval(%TuBigS) [[BIGS_ARG_ATTRS2:.*]])
#[no_mangle]
pub extern fn test_TuBigS(_: TuBigS) -> TuBigS { loop {} }

// CHECK: define void @test_TeBigS(%"TeBigS::Variant"* [[BIGS_RET_ATTRS]], %"TeBigS::Variant"* [[BIGS_ARG_ATTRS1]] byval(%"TeBigS::Variant") [[BIGS_ARG_ATTRS2]])
#[no_mangle]
pub extern fn test_TeBigS(_: TeBigS) -> TeBigS { loop {} }

#[derive(Clone, Copy)]
#[repr(C)]
pub union BigU {
    foo: [u32; 16],
}

#[repr(transparent)]
pub struct TsBigU(BigU);

#[repr(transparent)]
pub union TuBigU {
    field: BigU,
}

#[repr(transparent)]
pub enum TeBigU {
    Variant(BigU),
}

// CHECK: define void @test_BigU(%BigU* [[BIGU_RET_ATTRS:.*]], %BigU* [[BIGU_ARG_ATTRS1:.*]] byval(%BigU) [[BIGU_ARG_ATTRS2:.*]])
#[no_mangle]
pub extern fn test_BigU(_: BigU) -> BigU { loop {} }

// CHECK: define void @test_TsBigU(%TsBigU* [[BIGU_RET_ATTRS:.*]], %TsBigU* [[BIGU_ARG_ATTRS1]] byval(%TsBigU) [[BIGU_ARG_ATTRS2]])
#[no_mangle]
pub extern fn test_TsBigU(_: TsBigU) -> TsBigU { loop {} }

// CHECK: define void @test_TuBigU(%TuBigU* [[BIGU_RET_ATTRS]], %TuBigU* [[BIGU_ARG_ATTRS1]] byval(%TuBigU) [[BIGU_ARG_ATTRS2]])
#[no_mangle]
pub extern fn test_TuBigU(_: TuBigU) -> TuBigU { loop {} }

// CHECK: define void @test_TeBigU(%"TeBigU::Variant"* [[BIGU_RET_ATTRS]], %"TeBigU::Variant"* [[BIGU_ARG_ATTRS1]] byval(%"TeBigU::Variant") [[BIGU_ARG_ATTRS2]])
#[no_mangle]
pub extern fn test_TeBigU(_: TeBigU) -> TeBigU { loop {} }
<commit_msg>Rollup merge of #73926 - joaopaulocarreiro:github_rust-6, r=nikomatsakis<commit_after>// compile-flags: -C no-prepopulate-passes
// ignore-tidy-linelength

// min-system-llvm-version: 9.0
// ignore-arm
// ignore-aarch64
// ignore-mips
// ignore-mips64
// ignore-powerpc
// ignore-powerpc64
// ignore-riscv64 see codegen/riscv-abi
// ignore-windows
// See repr-transparent.rs

#![feature(transparent_unions)]

#![crate_type="lib"]

#[derive(Clone, Copy)]
#[repr(C)]
pub struct BigS([u32; 16]);

#[repr(transparent)]
pub struct TsBigS(BigS);

#[repr(transparent)]
pub union TuBigS {
    field: BigS,
}

#[repr(transparent)]
pub enum TeBigS {
    Variant(BigS),
}

// CHECK: define void @test_BigS(%BigS* [[BIGS_RET_ATTRS:.*]], %BigS* [[BIGS_ARG_ATTRS1:.*]] byval(%BigS) [[BIGS_ARG_ATTRS2:.*]])
#[no_mangle]
pub extern fn test_BigS(_: BigS) -> BigS { loop {} }

// CHECK: define void @test_TsBigS(%TsBigS* [[BIGS_RET_ATTRS]], %TsBigS* [[BIGS_ARG_ATTRS1]] byval(%TsBigS) [[BIGS_ARG_ATTRS2:.*]])
#[no_mangle]
pub extern fn test_TsBigS(_: TsBigS) -> TsBigS { loop {} }

// CHECK: define void @test_TuBigS(%TuBigS* [[BIGS_RET_ATTRS]], %TuBigS* [[BIGS_ARG_ATTRS1]] byval(%TuBigS) [[BIGS_ARG_ATTRS2:.*]])
#[no_mangle]
pub extern fn test_TuBigS(_: TuBigS) -> TuBigS { loop {} }

// CHECK: define void @test_TeBigS(%"TeBigS::Variant"* [[BIGS_RET_ATTRS]], %"TeBigS::Variant"* [[BIGS_ARG_ATTRS1]] byval(%"TeBigS::Variant") [[BIGS_ARG_ATTRS2]])
#[no_mangle]
pub extern fn test_TeBigS(_: TeBigS) -> TeBigS { loop {} }

#[derive(Clone, Copy)]
#[repr(C)]
pub union BigU {
    foo: [u32; 16],
}

#[repr(transparent)]
pub struct TsBigU(BigU);

#[repr(transparent)]
pub union TuBigU {
    field: BigU,
}

#[repr(transparent)]
pub enum TeBigU {
    Variant(BigU),
}

// CHECK: define void @test_BigU(%BigU* [[BIGU_RET_ATTRS:.*]], %BigU* [[BIGU_ARG_ATTRS1:.*]] byval(%BigU) [[BIGU_ARG_ATTRS2:.*]])
#[no_mangle]
pub extern fn test_BigU(_: BigU) -> BigU { loop {} }

// CHECK: define void @test_TsBigU(%TsBigU* [[BIGU_RET_ATTRS:.*]], %TsBigU* [[BIGU_ARG_ATTRS1]] byval(%TsBigU) [[BIGU_ARG_ATTRS2]])
#[no_mangle]
pub extern fn test_TsBigU(_: TsBigU) -> TsBigU { loop {} }

// CHECK: define void @test_TuBigU(%TuBigU* [[BIGU_RET_ATTRS]], %TuBigU* [[BIGU_ARG_ATTRS1]] byval(%TuBigU) [[BIGU_ARG_ATTRS2]])
#[no_mangle]
pub extern fn test_TuBigU(_: TuBigU) -> TuBigU { loop {} }

// CHECK: define void @test_TeBigU(%"TeBigU::Variant"* [[BIGU_RET_ATTRS]], %"TeBigU::Variant"* [[BIGU_ARG_ATTRS1]] byval(%"TeBigU::Variant") [[BIGU_ARG_ATTRS2]])
#[no_mangle]
pub extern fn test_TeBigU(_: TeBigU) -> TeBigU { loop {} }
<|endoftext|>


<commit_before>use amq_protocol::frame::{BackToTheBuffer, GenError, GenResult, WriteContext};
use std::{cmp, io};

#[derive(Debug, PartialEq, Clone)]
pub(crate) struct Buffer {
    memory: Vec<u8>,
    capacity: usize,
    position: usize,
    end: usize,
}

pub(crate) struct Checkpoint(usize);

impl Buffer {
    pub(crate) fn with_capacity(capacity: usize) -> Buffer {
        Buffer {
            memory: vec![0; capacity],
            capacity,
            position: 0,
            end: 0,
        }
    }

    pub(crate) fn checkpoint(&self) -> Checkpoint {
        Checkpoint(self.end)
    }

    pub(crate) fn rollback(&mut self, checkpoint: Checkpoint) {
        self.end = checkpoint.0;
    }

    pub(crate) fn grow(&mut self, new_size: usize) -> bool {
        if self.capacity >= new_size {
            return false;
        }

        self.memory.resize(new_size, 0);
        self.capacity = new_size;
        true
    }

    pub(crate) fn available_data(&self) -> usize {
        self.end - self.position
    }

    pub(crate) fn available_space(&self) -> usize {
        self.capacity - self.end
    }

    pub(crate) fn consume(&mut self, count: usize) -> usize {
        let cnt = cmp::min(count, self.available_data());
        self.position += cnt;
        cnt
    }

    pub(crate) fn fill(&mut self, count: usize) -> usize {
        let cnt = cmp::min(count, self.available_space());
        self.end += cnt;
        cnt
    }

    pub(crate) fn data(&self) -> &[u8] {
        &self.memory[self.position..self.end]
    }

    pub(crate) fn space(&mut self) -> &mut [u8] {
        &mut self.memory[self.end..self.capacity]
    }

    pub(crate) fn shift(&mut self) {
        let length = self.end - self.position;
        if length > self.position {
            return;
        } else {
            let (start, end) = self.memory.split_at_mut(self.position);
            start[..length].copy_from_slice(&end[..length]);
        }
        self.position = 0;
        self.end = length;
    }

    pub(crate) fn shift_unless_available(&mut self, size: usize) {
        if self.available_space() < size {
            self.shift();
        }
    }
}

impl io::Write for &mut Buffer {
    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
        let amt = self.space().write(data)?;
        self.fill(amt);
        Ok(amt)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl BackToTheBuffer for &mut Buffer {
    fn reserve_write_use<
        Tmp,
        Gen: Fn(WriteContext<Self>) -> Result<(WriteContext<Self>, Tmp), GenError>,
        Before: Fn(WriteContext<Self>, Tmp) -> GenResult<Self>,
    >(
        s: WriteContext<Self>,
        reserved: usize,
        gen: &Gen,
        before: &Before,
    ) -> Result<WriteContext<Self>, GenError> {
        if s.write.available_space() < reserved {
            return Err(GenError::BufferTooSmall(
                reserved - s.write.available_space(),
            ));
        }
        let start = s.write.checkpoint();
        s.write.fill(reserved);
        gen(s).and_then(|(s, tmp)| {
            let end = s.write.checkpoint();
            s.write.rollback(start);
            before(s, tmp).and_then(|s| {
                s.write.rollback(end);
                Ok(s)
            })
        })
    }
}

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

    #[test]
    fn fill_and_consume() {
        let mut b = Buffer::with_capacity(10);
        assert_eq!(b.available_data(), 0);
        assert_eq!(b.available_space(), 10);
        let res = b.space().write(&b"abcd"[..]).map(|size| {
            b.fill(size);
            size
        });
        assert_eq!(res.ok(), Some(4));
        assert_eq!(b.available_data(), 4);
        assert_eq!(b.available_space(), 6);

        assert_eq!(b.data(), &b"abcd"[..]);

        b.consume(2);
        assert_eq!(b.available_data(), 2);
        assert_eq!(b.available_space(), 6);
        assert_eq!(b.data(), &b"cd"[..]);

        b.shift();
        assert_eq!(b.available_data(), 2);
        assert_eq!(b.available_space(), 8);
        assert_eq!(b.data(), &b"cd"[..]);

        assert_eq!(
            b.space()
                .write(&b"efghijklmnop"[..])
                .map(|size| {
                    b.fill(size);
                    size
                })
                .ok(),
            Some(8)
        );
        assert_eq!(b.available_data(), 10);
        assert_eq!(b.available_space(), 0);
        assert_eq!(b.data(), &b"cdefghijkl"[..]);
        b.shift();
        assert_eq!(b.available_data(), 10);
        assert_eq!(b.available_space(), 0);
        assert_eq!(b.data(), &b"cdefghijkl"[..]);
    }
}
<commit_msg>buffer: drop old api tests<commit_after>use amq_protocol::frame::{BackToTheBuffer, GenError, GenResult, WriteContext};
use std::{cmp, io};

#[derive(Debug, PartialEq, Clone)]
pub(crate) struct Buffer {
    memory: Vec<u8>,
    capacity: usize,
    position: usize,
    end: usize,
}

pub(crate) struct Checkpoint(usize);

impl Buffer {
    pub(crate) fn with_capacity(capacity: usize) -> Buffer {
        Buffer {
            memory: vec![0; capacity],
            capacity,
            position: 0,
            end: 0,
        }
    }

    pub(crate) fn checkpoint(&self) -> Checkpoint {
        Checkpoint(self.end)
    }

    pub(crate) fn rollback(&mut self, checkpoint: Checkpoint) {
        self.end = checkpoint.0;
    }

    pub(crate) fn grow(&mut self, new_size: usize) -> bool {
        if self.capacity >= new_size {
            return false;
        }

        self.memory.resize(new_size, 0);
        self.capacity = new_size;
        true
    }

    pub(crate) fn available_data(&self) -> usize {
        self.end - self.position
    }

    pub(crate) fn available_space(&self) -> usize {
        self.capacity - self.end
    }

    pub(crate) fn consume(&mut self, count: usize) -> usize {
        let cnt = cmp::min(count, self.available_data());
        self.position += cnt;
        cnt
    }

    pub(crate) fn fill(&mut self, count: usize) -> usize {
        let cnt = cmp::min(count, self.available_space());
        self.end += cnt;
        cnt
    }

    pub(crate) fn data(&self) -> &[u8] {
        &self.memory[self.position..self.end]
    }

    pub(crate) fn space(&mut self) -> &mut [u8] {
        &mut self.memory[self.end..self.capacity]
    }

    pub(crate) fn shift(&mut self) {
        let length = self.end - self.position;
        if length > self.position {
            return;
        } else {
            let (start, end) = self.memory.split_at_mut(self.position);
            start[..length].copy_from_slice(&end[..length]);
        }
        self.position = 0;
        self.end = length;
    }

    pub(crate) fn shift_unless_available(&mut self, size: usize) {
        if self.available_space() < size {
            self.shift();
        }
    }
}

impl io::Write for &mut Buffer {
    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
        let amt = self.space().write(data)?;
        self.fill(amt);
        Ok(amt)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl BackToTheBuffer for &mut Buffer {
    fn reserve_write_use<
        Tmp,
        Gen: Fn(WriteContext<Self>) -> Result<(WriteContext<Self>, Tmp), GenError>,
        Before: Fn(WriteContext<Self>, Tmp) -> GenResult<Self>,
    >(
        s: WriteContext<Self>,
        reserved: usize,
        gen: &Gen,
        before: &Before,
    ) -> Result<WriteContext<Self>, GenError> {
        if s.write.available_space() < reserved {
            return Err(GenError::BufferTooSmall(
                reserved - s.write.available_space(),
            ));
        }
        let start = s.write.checkpoint();
        s.write.fill(reserved);
        gen(s).and_then(|(s, tmp)| {
            let end = s.write.checkpoint();
            s.write.rollback(start);
            before(s, tmp).and_then(|s| {
                s.write.rollback(end);
                Ok(s)
            })
        })
    }
}
<|endoftext|>


<commit_before>use amq_protocol::frame::{BackToTheBuffer, GenError, GenResult, WriteContext};
use std::{
    cmp,
    io::{self, IoSlice, IoSliceMut},
};

#[derive(Debug, PartialEq, Clone)]
pub(crate) struct Buffer {
    memory: Vec<u8>,
    capacity: usize,
    position: usize,
    end: usize,
}

pub(crate) struct Checkpoint(usize);

impl Buffer {
    pub(crate) fn with_capacity(capacity: usize) -> Buffer {
        Buffer {
            memory: vec![0; capacity],
            capacity,
            position: 0,
            end: 0,
        }
    }

    pub(crate) fn checkpoint(&self) -> Checkpoint {
        Checkpoint(self.end)
    }

    pub(crate) fn rollback(&mut self, checkpoint: Checkpoint) {
        self.end = checkpoint.0;
    }

    pub(crate) fn grow(&mut self, new_size: usize) -> bool {
        if self.capacity >= new_size {
            return false;
        }

        let old_capacity = self.capacity;

        self.memory.resize(new_size, 0);
        self.capacity = new_size;

        if self.end < self.position {
            let (old, new) = self.memory.split_at_mut(old_capacity);
            new[..].copy_from_slice(&old[..self.end]);
            self.end += old_capacity;
        }

        true
    }

    pub(crate) fn available_data(&self) -> usize {
        if self.end >= self.position {
            self.end - self.position
        } else {
            self.end + self.capacity - self.position
        }
    }

    fn available_space(&self) -> usize {
        if self.end >= self.position {
            self.capacity - self.end + self.position
        } else {
            self.position - self.end
        }
    }

    pub(crate) fn consume(&mut self, count: usize, ring: bool) -> usize {
        let cnt = cmp::min(count, self.available_data());
        self.position += cnt;
        if ring {
            self.position %= self.capacity;
        }
        cnt
    }

    pub(crate) fn fill(&mut self, count: usize, ring: bool) -> usize {
        let cnt = cmp::min(count, self.available_space());
        self.end += cnt;
        if ring {
            self.end %= self.capacity;
        }
        cnt
    }

    pub(crate) fn write_to<T: io::Write>(&self, writer: &mut T) -> io::Result<usize> {
        if self.end >= self.position {
            writer.write(&self.memory[self.position..self.end])
        } else {
            writer.write_vectored(&[
                IoSlice::new(&self.memory[self.position..]),
                IoSlice::new(&self.memory[..self.end]),
            ])
        }
    }

    pub(crate) fn read_from<T: io::Read>(
        &mut self,
        reader: &mut T,
        ring: bool,
    ) -> io::Result<usize> {
        if ring {
            if self.end >= self.position {
                let (start, end) = self.memory.split_at_mut(self.end);
                reader.read_vectored(
                    &mut [
                        IoSliceMut::new(&mut end[..]),
                        IoSliceMut::new(&mut start[..self.position]),
                    ][..],
                )
            } else {
                reader.read(&mut self.memory[self.end..self.position])
            }
        } else {
            reader.read(&mut self.memory[self.end..])
        }
    }

    pub(crate) fn offset(&self, buf: &[u8]) -> usize {
        let bufptr = buf.as_ptr() as usize;
        if self.end >= self.position {
            let data = &self.memory[self.position..self.end];
            bufptr - data.as_ptr() as usize
        } else {
            let data = &self.memory[self.position..];
            let dataptr = data.as_ptr() as usize;
            if dataptr < bufptr {
                bufptr - dataptr
            } else {
                let data = &self.memory[..self.end];
                bufptr + self.capacity - self.position - data.as_ptr() as usize
            }
        }
    }

    pub(crate) fn data(&self) -> &[u8] {
        // FIXME: drop
        &self.memory[self.position..self.end]
    }

    pub(crate) fn shift(&mut self) {
        // FIXME: drop
        let length = self.end - self.position;
        if length > self.position {
            return;
        } else {
            let (start, end) = self.memory.split_at_mut(self.position);
            start[..length].copy_from_slice(&end[..length]);
        }
        self.position = 0;
        self.end = length;
    }

    pub(crate) fn shift_unless_available(&mut self, size: usize) {
        // FIXME: drop
        if self.available_space() < size {
            self.shift();
        }
    }
}

impl io::Write for &mut Buffer {
    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
        let amt = if self.end >= self.position {
            let mut space = &mut self.memory[self.end..];
            let mut amt = space.write(data)?;
            if amt == self.capacity - self.end {
                let mut space = &mut self.memory[..self.position];
                amt += space.write(&data[amt..])?;
            }
            amt
        } else {
            let mut space = &mut self.memory[self.end..self.position];
            space.write(data)?
        };
        self.fill(amt, true);
        Ok(amt)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl BackToTheBuffer for &mut Buffer {
    fn reserve_write_use<
        Tmp,
        Gen: Fn(WriteContext<Self>) -> Result<(WriteContext<Self>, Tmp), GenError>,
        Before: Fn(WriteContext<Self>, Tmp) -> GenResult<Self>,
    >(
        s: WriteContext<Self>,
        reserved: usize,
        gen: &Gen,
        before: &Before,
    ) -> Result<WriteContext<Self>, GenError> {
        if s.write.available_space() < reserved {
            return Err(GenError::BufferTooSmall(
                reserved - s.write.available_space(),
            ));
        }
        let start = s.write.checkpoint();
        s.write.fill(reserved, true);
        gen(s).and_then(|(s, tmp)| {
            let end = s.write.checkpoint();
            s.write.rollback(start);
            before(s, tmp).and_then(|s| {
                s.write.rollback(end);
                Ok(s)
            })
        })
    }
}
<commit_msg>buffer: fix edge cases and add more safe guards<commit_after>use amq_protocol::frame::{BackToTheBuffer, GenError, GenResult, WriteContext};
use std::{
    cmp,
    io::{self, IoSlice, IoSliceMut},
};

#[derive(Debug, PartialEq, Clone)]
pub(crate) struct Buffer {
    memory: Vec<u8>,
    capacity: usize,
    position: usize,
    end: usize,
    available_data: usize,
}

// The bool is true when the checkpoint is to go backwards, false for forward
pub(crate) struct Checkpoint(usize, bool);

impl Buffer {
    pub(crate) fn with_capacity(capacity: usize) -> Buffer {
        Buffer {
            memory: vec![0; capacity],
            capacity,
            position: 0,
            end: 0,
            available_data: 0,
        }
    }

    pub(crate) fn checkpoint(&self) -> Checkpoint {
        Checkpoint(self.end, true)
    }

    pub(crate) fn rollback(&mut self, checkpoint: Checkpoint) {
        if checkpoint.1 {
            if self.end > checkpoint.0 {
                self.available_data -= self.end - checkpoint.0;
            } else {
                self.available_data -= self.end + (self.capacity - checkpoint.0);
            }
        } else {
            if self.end > checkpoint.0 {
                self.available_data += (self.capacity - self.end) + checkpoint.0;
            } else {
                self.available_data += checkpoint.0 - self.end;
            }
        }
        self.end = checkpoint.0;
    }

    pub(crate) fn grow(&mut self, new_size: usize) -> bool {
        if self.capacity >= new_size {
            return false;
        }

        let old_capacity = self.capacity;

        self.memory.resize(new_size, 0);
        self.capacity = new_size;

        if self.end < self.position {
            // FIXME: make sure there's enough room
            let (old, new) = self.memory.split_at_mut(old_capacity);
            new[..].copy_from_slice(&old[..self.end]);
            self.end += old_capacity;
        }

        true
    }

    pub(crate) fn available_data(&self) -> usize {
        self.available_data
    }

    fn available_space(&self) -> usize {
        self.capacity - self.available_data
    }

    pub(crate) fn consume(&mut self, count: usize, ring: bool) -> usize {
        let cnt = cmp::min(count, self.available_data());
        self.position += cnt;
        if ring {
            self.position %= self.capacity;
        }
        self.available_data -= cnt;
        cnt
    }

    pub(crate) fn fill(&mut self, count: usize, ring: bool) -> usize {
        let cnt = cmp::min(count, self.available_space());
        self.end += cnt;
        if ring {
            self.end %= self.capacity;
        }
        self.available_data += cnt;
        cnt
    }

    pub(crate) fn write_to<T: io::Write>(&self, writer: &mut T) -> io::Result<usize> {
        if self.available_data() == 0 {
            Ok(0)
        } else {
            if self.end > self.position {
                writer.write(&self.memory[self.position..self.end])
            } else {
                writer.write_vectored(&[
                    IoSlice::new(&self.memory[self.position..]),
                    IoSlice::new(&self.memory[..self.end]),
                ])
            }
        }
    }

    pub(crate) fn read_from<T: io::Read>(&mut self, reader: &mut T, ring: bool) -> io::Result<usize> {
        if self.available_space() == 0 {
            Ok(0)
        } else {
            if ring {
                if self.end >= self.position {
                    let (start, end) = self.memory.split_at_mut(self.end);
                    reader.read_vectored(
                        &mut [
                        IoSliceMut::new(&mut end[..]),
                        IoSliceMut::new(&mut start[..self.position]),
                        ][..],
                    )
                } else {
                    reader.read(&mut self.memory[self.end..self.position])
                }
            } else {
                reader.read(&mut self.memory[self.end..])
            }
        }
    }

    pub(crate) fn offset(&self, buf: &[u8]) -> usize {
        let bufptr = buf.as_ptr() as usize;
        if self.end >= self.position {
            let data = &self.memory[self.position..self.end];
            bufptr - data.as_ptr() as usize
        } else {
            let data = &self.memory[self.position..];
            let dataptr = data.as_ptr() as usize;
            if dataptr < bufptr {
                bufptr - dataptr
            } else {
                let data = &self.memory[..self.end];
                bufptr + self.capacity - self.position - data.as_ptr() as usize
            }
        }
    }

    pub(crate) fn data(&self) -> &[u8] {
        // FIXME: drop
        &self.memory[self.position..self.end]
    }

    pub(crate) fn shift(&mut self) {
        // FIXME: drop
        let length = self.end - self.position;
        if length > self.position {
            return;
        } else {
            let (start, end) = self.memory.split_at_mut(self.position);
            start[..length].copy_from_slice(&end[..length]);
        }
        self.position = 0;
        self.end = length;
    }

    pub(crate) fn shift_unless_available(&mut self, size: usize) {
        // FIXME: drop
        if self.available_space() < size {
            self.shift();
        }
    }
}

impl io::Write for &mut Buffer {
    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
        let amt = if self.available_space() == 0 {
            0
        } else {
            if self.end >= self.position {
                let mut space = &mut self.memory[self.end..];
                let mut amt = space.write(data)?;
                if amt == self.capacity - self.end {
                    let mut space = &mut self.memory[..self.position];
                    amt += space.write(&data[amt..])?;
                }
                amt
            } else {
                let mut space = &mut self.memory[self.end..self.position];
                space.write(data)?
            }
        };
        self.fill(amt, true);
        Ok(amt)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl BackToTheBuffer for &mut Buffer {
    fn reserve_write_use<
        Tmp,
        Gen: Fn(WriteContext<Self>) -> Result<(WriteContext<Self>, Tmp), GenError>,
        Before: Fn(WriteContext<Self>, Tmp) -> GenResult<Self>,
    >(
        s: WriteContext<Self>,
        reserved: usize,
        gen: &Gen,
        before: &Before,
    ) -> Result<WriteContext<Self>, GenError> {
        if s.write.available_space() < reserved {
            return Err(GenError::BufferTooSmall(
                reserved - s.write.available_space(),
            ));
        }
        let start = s.write.checkpoint();
        s.write.fill(reserved, true);
        gen(s).and_then(|(s, tmp)| {
            let mut end = s.write.checkpoint();
            end.1 = false;
            s.write.rollback(start);
            before(s, tmp).and_then(|s| {
                s.write.rollback(end);
                Ok(s)
            })
        })
    }
}
<|endoftext|>


<commit_before>use collections::HashMap;
use std::io::{Open, Read, Write, TempDir};
use std::num::Zero;
use std::comm::{Data, Empty, Disconnected};
use sync::{Arc, RWLock};
use sync::Future;
use rand::task_rng;
use rand::distributions::{IndependentSample, Range};
use rand::distributions::range::SampleRange;
use native::io::file;
use native::io::file::FileDesc;
use std::rt::rtio::RtioFileStream;

/* Linux seems to use 4K segments, that's a good bet */
pub static PAGE_SIZE: uint = 4 * 1024;
/*
 * how much of the page_id should be reserved for pages?
 * the rest goes for segments. The bigger this is, the larger the segment files are
 * and the more pages they contain. (2^PAGE_BITS)*PAGE_SIZE gives you segment size.
 *
 * Adjust as required.
 */
pub static PAGE_BITS: uint = 32;

pub type ConcurrentFrame = Arc<RWLock<BufferFrame>>;

pub struct BufferManager {
	size: uint,
	entries: HashMap<u64, BufferEntry>,
	path: Path,
}

struct BufferEntry {
	frame: ConcurrentFrame,
	written: Cleanliness,
}

#[deriving(Eq)]
enum Status {
	Free,
	Fixed(uint),
}

#[deriving(Eq)]
enum Cleanliness {
	Clean,
	Dirty,
}

pub struct BufferFrame {
	page_id: u64,
	data: Vec<u8>,
	fixed: Status,
}

/* Destructor trait implementation */
impl Drop for BufferManager {
	fn drop(&mut self) {
		for (page, entry) in self.entries.iter() {
			info!("Drop k == {}, dirty? {:?}", page, entry.written);
			match entry.written {
				Clean => (),
				Dirty => {
					let frame = entry.frame.read();
					self.write_page(frame.page_id, frame.get_data());
				},
			}
		}
	}
}

impl BufferManager {
	pub fn new(size: uint, path: Path) -> BufferManager {
		let h = HashMap::with_capacity(size);
		BufferManager {size: size, entries: h, path: path}
	}

	fn open_or_create(&self, page_id: u64) -> FileDesc {
		let (segment, _) = split_segment(page_id);
		let file_path = self.path.join(segment.to_str());
		match file::open(&file_path.to_c_str(), Open, Read) {
			Ok(f) => return f,
			Err(_) => self.create(page_id),
		};
		match file::open(&file_path.to_c_str(), Open, Read) {
			Ok(f) => f,
			Err(e) => fail!("Failed reopening written file: {}", e),
		}
	}

	fn create(&self, page_id: u64) {
		let (segment, offset) = split_segment(page_id);
		let path = self.path.join(segment.to_str());

		match file::open(&path.to_c_str(), Open, Write) {
			Ok(mut f) => {
				match f.pwrite([0_u8, ..PAGE_SIZE], offset * PAGE_SIZE as u64) {
					Ok(_) => (),
					//Err(e) => fail!("Failed pwriting to file {}: {}", &path.to_c_str().as_str(), e),
					// in this case, err means everything is fine, since there is a bug in Rust' libnative
					Err(_) => (),
				};
			},
			Err(e) => fail!("Failed opening file for write: {}", e),
		};
	}

	/*
	 * returns true if page could be loaded
	 */
	fn load_page(&mut self, page_id: u64) -> bool {
		if self.entries.len() == self.size {
			if !self.evict_page() {
				return false;
			}
		}
		let (_, offset) = split_segment(page_id);

		let mut file_handle = self.open_or_create(page_id);
		let mut buf = [0_u8, ..PAGE_SIZE];
		let content = match file_handle.pread(buf, offset * PAGE_SIZE as u64) {
			Ok(_) => {
				Vec::from_slice(buf)
				/*
				if n == PAGE_SIZE as int {
					Vec::from_slice(buf)
				} else {
					fail!("pread failed: wanted {}, got {}",
						PAGE_SIZE, n)
				}*/
			},
			Err(e) => fail!("Couldn't read from page: {}", e),
		};

		let frame = BufferFrame {data: content, page_id: page_id, fixed: Free};
		let entry = BufferEntry {frame: Arc::new(RWLock::new(frame)), written: Clean};
		self.entries.insert(page_id, entry);
		true
	}

	/*
	 * return false if no page could be evicted
	 */
	fn evict_page(&mut self) -> bool {
		let random_key = {
			let mut free_keys = self.entries.iter().filter_map(|e| {
				// iterate over entries, check if they are free
				// and return the keys of the free entries
				let (k,v) = e;
				let frame = v.frame.read();
				if frame.fixed == Free {
					Some(k)
				} else {
					None
				}
			});
			// pick a random key
			sample(&mut free_keys).map(|v| v.clone())
		};

		match random_key {
			None => false,
			Some(key) => {
				match self.entries.pop(&key) {
					None => false,
					Some(entry) => {
						info!("Evicting entry: {:?}", entry);
						if entry.written == Dirty {
							let frame = entry.frame.read();
							self.write_page(frame.page_id, frame.get_data());
						}
						true
					},
				}
			},
		}
	}

	pub fn fix_page(&mut self, page_id: u64) -> Option<ConcurrentFrame> {
		if !self.entries.contains_key(&page_id) {
			if !self.load_page(page_id) {
				return None;
			}
		}
		let entry = self.entries.get(&page_id);
		{
			let mut f = entry.frame.write();
			f.fixed = match f.fixed {
				Free => Fixed(1),
				Fixed(n) => Fixed(n+1),
			};
		}
		// Arcs can be cloned and they will all point to the same RWLock
		Some(entry.frame.clone())
	}

	pub fn unfix_page(&mut self, frame: ConcurrentFrame, is_dirty: bool) {
		{
			let mut frame = frame.write();
			frame.fixed = match frame.fixed {
				Fixed(1) => Free,
				Fixed(n) => Fixed(n-1),
				Free => fail!("Unfixing unfixed page"),
			};
		}

		if is_dirty {
			let frame = frame.read();
			let entry = self.entries.get_mut(&frame.page_id);
			entry.written = Dirty;
		}
	}

	fn write_page(&self, page_id: u64, data: &[u8]) {
		let (segment, offset) = split_segment(page_id);
		let file_path = self.path.join(segment.to_str());
		let mut handle = match file::open(&file_path.to_c_str(), Open, Write) {
			Ok(handle) => handle,
			Err(e) => fail!("Opening file for writing failed: {}", e),
		};

		info!("Writing to segment {}, offset {}", segment, offset);
		match handle.pwrite(data, offset * PAGE_SIZE as u64) {
			Ok(_) => (),
			//Err(e) => fail!("Writing to file failed: {}", e),
			// again Rust 0.10 pwrite error
			Err(_) => (),
		};
	}
}

/*
 * A buffer frame, the unit that gets the data
 */
impl BufferFrame {
	/*
	 * get_mut_data returns a borrowed reference to a mutable slice to the vector
	 * contents.
	 */
	pub fn get_mut_data<'a>(&'a mut self) -> &'a mut [u8] {
		self.data.as_mut_slice()
	}
	pub fn get_data<'a>(&'a self) -> &'a [u8] {
		self.data.as_slice()
	}
}

fn sample<'a, T, I:Iterator<T>>(from: &'a mut I) -> Option<T> {
	let from: ~[T] = from.collect();
	let l = from.len();
	if l == 0 {
		return None;
	}
	let index = randrange(l);
	Some(from[index])
}

fn randrange<X: SampleRange + Ord + Zero>(high: X) -> X {
	let between: Range<X> = Range::new(Zero::zero(), high);
	let mut rng = task_rng();
	between.ind_sample(&mut rng)
}

fn split_segment(num: u64) -> (u64, u64) {
	let high = num >> PAGE_BITS;
	let low = num & ((1 << PAGE_BITS) - 1);
	(high, low)
}

pub fn join_segment(segment: u64, page: u64) -> u64 {
	// TODO do proper masking before
	let high = segment << PAGE_BITS;
	let low = page;
	high & low
}

#[test]
fn test_create() {
	let dir = match TempDir::new("buffermanager") {
		Some(temp_dir) => temp_dir,
		None => fail!("creation of temporary directory"),
	};

	let mut bm = BufferManager::new(16, dir.path().clone());
	let pageref = match bm.fix_page(42) {
		Some(p) => p,
		None => fail!("Getting page failed"),
	};
	{
		let mut page = pageref.write();
		let data = page.get_mut_data();
		data[0] = 42;
	}
	bm.unfix_page(pageref, true);
}

#[test]
fn test_threads() {
	use rand::random;
	use std::task::spawn;
	use std::os;

	let pages_in_ram: uint = match os::getenv("PAGES_IN_RAM") {
		Some(v) => from_str(v).expect("PAGES_IN_RAM expects integer"),
		None => 1,
	};
	let pages_on_disk: u64 = match os::getenv("PAGES_ON_DISK") {
		Some(v) => from_str(v).expect("PAGES_ON_DISK expects integer"),
		None => 20,
	};
	let thread_count: uint = match os::getenv("THREADS") {
		Some(v) => from_str(v).expect("THREADS expects integer"),
		None => 3,
	};
	let dir = match TempDir::new("buffermanager") {
		Some(temp_dir) => temp_dir,
		None => fail!("creation of temporary directory"),
	};
	let p = dir.path();
	//let p = Path::new(".");

	let mut buffermanager = BufferManager::new(pages_in_ram, p.clone());

	for i in range(0, pages_on_disk) {
		let bf = match buffermanager.fix_page(i) {
			Some(frame) => frame,
			None => fail!("Couldn't fix page {}", i),
		};
		{
			let mut lock = bf.write();
			lock.get_mut_data()[0] = 0;
		}
		buffermanager.unfix_page(bf, true);
	}
	let bm = Arc::new(RWLock::new(buffermanager));

	// start scan thread
	let (tx, rx) = channel();
	let bm_scan = bm.clone();
	let mut scan = Future::spawn(proc() {
		let mut counters = Vec::from_elem(pages_on_disk as uint, 0_u8);
		loop {
			info!("Running scan thread...");
			match rx.try_recv() {
				Empty => {
					let page_number = randrange(pages_on_disk);
					let mut bm = bm_scan.write();
					let bf = match bm.fix_page(page_number) {
						Some(frame) => frame,
						None => fail!("Couldn't scan/fix page"),
					};
					let current_val = {
						let lock = bf.read();
						lock.get_data()[0]
					};
					bm.unfix_page(bf, false);
					bm.downgrade();
					// check if the value is going up
					assert!(&current_val >= counters.get(page_number as uint));
					// set to the new value
					let v = counters.get_mut(page_number as uint);
					*v = current_val;

				},
				/*
				 * terminate on both disconnect and when any kind of
				 * data arrives
				 */
				Disconnected => break,
				Data(_) => break,
			};
		}
	});

	// start read/write threads
	let mut rw_threads: Vec<Future<int>> = Vec::new();
	for _ in range(0, thread_count) {
		let bm = bm.clone();
		rw_threads.push(Future::spawn(proc() {
			let is_write = random::<bool>();
			info!("Creating new {} task",
				if is_write {"write"} else {"read"});
			let page_number = randrange(pages_on_disk);
			let mut bm = bm.write();
			if is_write {
				let bf = match bm.fix_page(page_number) {
					Some(frame) => frame,
					None => fail!("Couldn't fix page"),
				};
				{
					let mut lock = bf.write();
					let data = lock.get_mut_data();
					data[0] = data[0] + 1;
					info!("Wrote to page {}", page_number);
					debug!("data: {}", Vec::from_slice(data));
				}
				bm.unfix_page(bf, is_write);
			} else {
				let bf = match bm.fix_page(page_number) {
					Some(frame) => frame,
					None => fail!("Couldn't fix page"),
				};
				bm.unfix_page(bf, is_write);
			}
			// return whether we wrote (1) or read (0) as future
			if is_write {1} else {0}
		}));
	}

	// Rust does not have join, but we can wait on Futures which does the same
	let total_count = rw_threads.mut_iter().fold(0, |acc, val| acc + val.get());
	info!("Total count in RAM: {}", total_count);

	// terminate the scan thread and wait until it has completed
	tx.send("terminate");
	scan.get();

	// re-open the pages and check whether all numbers got saved
	let mut bm = BufferManager::new(pages_in_ram, p.clone());
	let mut total_count_on_disk = 0;
	for i in range(0, pages_on_disk) {
		let bf = match bm.fix_page(i) {
			Some(frame) => frame,
			None => fail!("Couldn't fix page")
		};
		let value = {
			let lock = bf.read();
			let data = lock.get_data();
			data[0]
		};
		bm.unfix_page(bf, false);
		// cast up from u8 to int
		total_count_on_disk += value as int;
	}
	info!("Total count on disk: {}", total_count_on_disk);
	assert_eq!(total_count, total_count_on_disk);
}
<commit_msg>Fixed segment & page joining<commit_after>use collections::HashMap;
use std::io::{Open, Read, Write, TempDir};
use std::num::Zero;
use std::comm::{Data, Empty, Disconnected};
use sync::{Arc, RWLock};
use sync::Future;
use rand::task_rng;
use rand::distributions::{IndependentSample, Range};
use rand::distributions::range::SampleRange;
use native::io::file;
use native::io::file::FileDesc;
use std::rt::rtio::RtioFileStream;

/* Linux seems to use 4K segments, that's a good bet */
pub static PAGE_SIZE: uint = 4 * 1024;
/*
 * how much of the page_id should be reserved for pages?
 * the rest goes for segments. The bigger this is, the larger the segment files are
 * and the more pages they contain. (2^PAGE_BITS)*PAGE_SIZE gives you segment size.
 *
 * Adjust as required.
 */
pub static PAGE_BITS: uint = 32;

pub type ConcurrentFrame = Arc<RWLock<BufferFrame>>;

pub struct BufferManager {
	size: uint,
	entries: HashMap<u64, BufferEntry>,
	path: Path,
}

struct BufferEntry {
	frame: ConcurrentFrame,
	written: Cleanliness,
}

#[deriving(Eq)]
enum Status {
	Free,
	Fixed(uint),
}

#[deriving(Eq)]
enum Cleanliness {
	Clean,
	Dirty,
}

pub struct BufferFrame {
	page_id: u64,
	data: Vec<u8>,
	fixed: Status,
}

/* Destructor trait implementation */
impl Drop for BufferManager {
	fn drop(&mut self) {
		for (page, entry) in self.entries.iter() {
			info!("Drop k == {}, dirty? {:?}", page, entry.written);
			match entry.written {
				Clean => (),
				Dirty => {
					let frame = entry.frame.read();
					self.write_page(frame.page_id, frame.get_data());
				},
			}
		}
	}
}

impl BufferManager {
	pub fn new(size: uint, path: Path) -> BufferManager {
		let h = HashMap::with_capacity(size);
		BufferManager {size: size, entries: h, path: path}
	}

	fn open_or_create(&self, page_id: u64) -> FileDesc {
		let (segment, _) = split_segment(page_id);
		let file_path = self.path.join(segment.to_str());
		match file::open(&file_path.to_c_str(), Open, Read) {
			Ok(f) => return f,
			Err(_) => self.create(page_id),
		};
		match file::open(&file_path.to_c_str(), Open, Read) {
			Ok(f) => f,
			Err(e) => fail!("Failed reopening written file: {}", e),
		}
	}

	fn create(&self, page_id: u64) {
		let (segment, offset) = split_segment(page_id);
		let path = self.path.join(segment.to_str());

		match file::open(&path.to_c_str(), Open, Write) {
			Ok(mut f) => {
				match f.pwrite([0_u8, ..PAGE_SIZE], offset * PAGE_SIZE as u64) {
					Ok(_) => (),
					//Err(e) => fail!("Failed pwriting to file {}: {}", &path.to_c_str().as_str(), e),
					// in this case, err means everything is fine, since there is a bug in Rust' libnative
					Err(_) => (),
				};
			},
			Err(e) => fail!("Failed opening file for write: {}", e),
		};
	}

	/*
	 * returns true if page could be loaded
	 */
	fn load_page(&mut self, page_id: u64) -> bool {
		if self.entries.len() == self.size {
			if !self.evict_page() {
				return false;
			}
		}
		let (_, offset) = split_segment(page_id);

		let mut file_handle = self.open_or_create(page_id);
		let mut buf = [0_u8, ..PAGE_SIZE];
		let content = match file_handle.pread(buf, offset * PAGE_SIZE as u64) {
			Ok(_) => {
				Vec::from_slice(buf)
				/*
				if n == PAGE_SIZE as int {
					Vec::from_slice(buf)
				} else {
					fail!("pread failed: wanted {}, got {}",
						PAGE_SIZE, n)
				}*/
			},
			Err(e) => fail!("Couldn't read from page: {}", e),
		};

		let frame = BufferFrame {data: content, page_id: page_id, fixed: Free};
		let entry = BufferEntry {frame: Arc::new(RWLock::new(frame)), written: Clean};
		self.entries.insert(page_id, entry);
		true
	}

	/*
	 * return false if no page could be evicted
	 */
	fn evict_page(&mut self) -> bool {
		let random_key = {
			let mut free_keys = self.entries.iter().filter_map(|e| {
				// iterate over entries, check if they are free
				// and return the keys of the free entries
				let (k,v) = e;
				let frame = v.frame.read();
				if frame.fixed == Free {
					Some(k)
				} else {
					None
				}
			});
			// pick a random key
			sample(&mut free_keys).map(|v| v.clone())
		};

		match random_key {
			None => false,
			Some(key) => {
				match self.entries.pop(&key) {
					None => false,
					Some(entry) => {
						info!("Evicting entry: {:?}", entry);
						if entry.written == Dirty {
							let frame = entry.frame.read();
							self.write_page(frame.page_id, frame.get_data());
						}
						true
					},
				}
			},
		}
	}

	pub fn fix_page(&mut self, page_id: u64) -> Option<ConcurrentFrame> {
		if !self.entries.contains_key(&page_id) {
			if !self.load_page(page_id) {
				return None;
			}
		}
		let entry = self.entries.get(&page_id);
		{
			let mut f = entry.frame.write();
			f.fixed = match f.fixed {
				Free => Fixed(1),
				Fixed(n) => Fixed(n+1),
			};
		}
		// Arcs can be cloned and they will all point to the same RWLock
		Some(entry.frame.clone())
	}

	pub fn unfix_page(&mut self, frame: ConcurrentFrame, is_dirty: bool) {
		{
			let mut frame = frame.write();
			frame.fixed = match frame.fixed {
				Fixed(1) => Free,
				Fixed(n) => Fixed(n-1),
				Free => fail!("Unfixing unfixed page"),
			};
		}

		if is_dirty {
			let frame = frame.read();
			let entry = self.entries.get_mut(&frame.page_id);
			entry.written = Dirty;
		}
	}

	fn write_page(&self, page_id: u64, data: &[u8]) {
		let (segment, offset) = split_segment(page_id);
		let file_path = self.path.join(segment.to_str());
		let mut handle = match file::open(&file_path.to_c_str(), Open, Write) {
			Ok(handle) => handle,
			Err(e) => fail!("Opening file for writing failed: {}", e),
		};

		info!("Writing to segment {}, offset {}", segment, offset);
		match handle.pwrite(data, offset * PAGE_SIZE as u64) {
			Ok(_) => (),
			//Err(e) => fail!("Writing to file failed: {}", e),
			// again Rust 0.10 pwrite error
			Err(_) => (),
		};
	}
}

/*
 * A buffer frame, the unit that gets the data
 */
impl BufferFrame {
	/*
	 * get_mut_data returns a borrowed reference to a mutable slice to the vector
	 * contents.
	 */
	pub fn get_mut_data<'a>(&'a mut self) -> &'a mut [u8] {
		self.data.as_mut_slice()
	}
	pub fn get_data<'a>(&'a self) -> &'a [u8] {
		self.data.as_slice()
	}
}

fn sample<'a, T, I:Iterator<T>>(from: &'a mut I) -> Option<T> {
	let from: ~[T] = from.collect();
	let l = from.len();
	if l == 0 {
		return None;
	}
	let index = randrange(l);
	Some(from[index])
}

fn randrange<X: SampleRange + Ord + Zero>(high: X) -> X {
	let between: Range<X> = Range::new(Zero::zero(), high);
	let mut rng = task_rng();
	between.ind_sample(&mut rng)
}

fn split_segment(num: u64) -> (u64, u64) {
	let high = num >> PAGE_BITS;
	let low = num & ((1 << PAGE_BITS) - 1);
	(high, low)
}

pub fn join_segment(segment: u64, page: u64) -> u64 {
	// TODO do proper masking before
	let high = segment << PAGE_BITS;
	let low = page;
	println!("Joining {} with {} into {}", segment, page, high ^ low);
	high ^ low
}

#[test]
fn test_create() {
	let dir = match TempDir::new("buffermanager") {
		Some(temp_dir) => temp_dir,
		None => fail!("creation of temporary directory"),
	};

	let mut bm = BufferManager::new(16, dir.path().clone());
	let pageref = match bm.fix_page(42) {
		Some(p) => p,
		None => fail!("Getting page failed"),
	};
	{
		let mut page = pageref.write();
		let data = page.get_mut_data();
		data[0] = 42;
	}
	bm.unfix_page(pageref, true);
}

#[test]
fn test_threads() {
	use rand::random;
	use std::task::spawn;
	use std::os;

	let pages_in_ram: uint = match os::getenv("PAGES_IN_RAM") {
		Some(v) => from_str(v).expect("PAGES_IN_RAM expects integer"),
		None => 1,
	};
	let pages_on_disk: u64 = match os::getenv("PAGES_ON_DISK") {
		Some(v) => from_str(v).expect("PAGES_ON_DISK expects integer"),
		None => 20,
	};
	let thread_count: uint = match os::getenv("THREADS") {
		Some(v) => from_str(v).expect("THREADS expects integer"),
		None => 3,
	};
	let dir = match TempDir::new("buffermanager") {
		Some(temp_dir) => temp_dir,
		None => fail!("creation of temporary directory"),
	};
	let p = dir.path();
	//let p = Path::new(".");

	let mut buffermanager = BufferManager::new(pages_in_ram, p.clone());

	for i in range(0, pages_on_disk) {
		let bf = match buffermanager.fix_page(i) {
			Some(frame) => frame,
			None => fail!("Couldn't fix page {}", i),
		};
		{
			let mut lock = bf.write();
			lock.get_mut_data()[0] = 0;
		}
		buffermanager.unfix_page(bf, true);
	}
	let bm = Arc::new(RWLock::new(buffermanager));

	// start scan thread
	let (tx, rx) = channel();
	let bm_scan = bm.clone();
	let mut scan = Future::spawn(proc() {
		let mut counters = Vec::from_elem(pages_on_disk as uint, 0_u8);
		loop {
			info!("Running scan thread...");
			match rx.try_recv() {
				Empty => {
					let page_number = randrange(pages_on_disk);
					let mut bm = bm_scan.write();
					let bf = match bm.fix_page(page_number) {
						Some(frame) => frame,
						None => fail!("Couldn't scan/fix page"),
					};
					let current_val = {
						let lock = bf.read();
						lock.get_data()[0]
					};
					bm.unfix_page(bf, false);
					bm.downgrade();
					// check if the value is going up
					assert!(&current_val >= counters.get(page_number as uint));
					// set to the new value
					let v = counters.get_mut(page_number as uint);
					*v = current_val;

				},
				/*
				 * terminate on both disconnect and when any kind of
				 * data arrives
				 */
				Disconnected => break,
				Data(_) => break,
			};
		}
	});

	// start read/write threads
	let mut rw_threads: Vec<Future<int>> = Vec::new();
	for _ in range(0, thread_count) {
		let bm = bm.clone();
		rw_threads.push(Future::spawn(proc() {
			let is_write = random::<bool>();
			info!("Creating new {} task",
				if is_write {"write"} else {"read"});
			let page_number = randrange(pages_on_disk);
			let mut bm = bm.write();
			if is_write {
				let bf = match bm.fix_page(page_number) {
					Some(frame) => frame,
					None => fail!("Couldn't fix page"),
				};
				{
					let mut lock = bf.write();
					let data = lock.get_mut_data();
					data[0] = data[0] + 1;
					info!("Wrote to page {}", page_number);
					debug!("data: {}", Vec::from_slice(data));
				}
				bm.unfix_page(bf, is_write);
			} else {
				let bf = match bm.fix_page(page_number) {
					Some(frame) => frame,
					None => fail!("Couldn't fix page"),
				};
				bm.unfix_page(bf, is_write);
			}
			// return whether we wrote (1) or read (0) as future
			if is_write {1} else {0}
		}));
	}

	// Rust does not have join, but we can wait on Futures which does the same
	let total_count = rw_threads.mut_iter().fold(0, |acc, val| acc + val.get());
	info!("Total count in RAM: {}", total_count);

	// terminate the scan thread and wait until it has completed
	tx.send("terminate");
	scan.get();

	// re-open the pages and check whether all numbers got saved
	let mut bm = BufferManager::new(pages_in_ram, p.clone());
	let mut total_count_on_disk = 0;
	for i in range(0, pages_on_disk) {
		let bf = match bm.fix_page(i) {
			Some(frame) => frame,
			None => fail!("Couldn't fix page")
		};
		let value = {
			let lock = bf.read();
			let data = lock.get_data();
			data[0]
		};
		bm.unfix_page(bf, false);
		// cast up from u8 to int
		total_count_on_disk += value as int;
	}
	info!("Total count on disk: {}", total_count_on_disk);
	assert_eq!(total_count, total_count_on_disk);
}
<|endoftext|>